/* -*- Mode: C; tab-width: 2; indent-tabs-mode: nil; c-basic-offset: 2 -*- */
/* vim:set sw=2 sts=2 et cin: */
/*
 * This file is part of the MUSE Instrument Pipeline
 * Copyright (C) 2017-2018 European Southern Observatory
 *
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 2 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, write to the Free Software
 * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.
 */

#include <string.h>

#include "muse_raman.h"
#include <muse.h>

/*----------------------------------------------------------------------------*/
/**
 * @defgroup muse_raman        Raman scattering light removal procedures
 *
 * This module contains functionality to remove the Raman scattering lines
 * of the O2 and N2 molecules, exited around the Laser Guide Stars of the
 * 4LGSF used by the GALACSI AO system of MUSE.
 */
/*----------------------------------------------------------------------------*/

/**@{*/

/*----------------------------------------------------------------------------*/
/**
   @brief  Load the raman lines file.
   @param  aProcessing   the processing structure
   @return Table with sky lines
   @cpl_ensure{aFile, CPL_ERROR_NULL_INPUT, NULL}
 */
/*----------------------------------------------------------------------------*/
cpl_table *
muse_raman_lines_load(muse_processing *aProcessing)
{
  cpl_ensure(aProcessing, CPL_ERROR_NULL_INPUT, NULL);
  cpl_frameset *frames = muse_frameset_find(aProcessing->inframes,
                                            MUSE_TAG_RAMAN_LINES, 0, CPL_FALSE);
  cpl_errorstate es = cpl_errorstate_get();
  cpl_frame *frame = cpl_frameset_get_position(frames, 0);
  if (!frame) {
    cpl_frameset_delete(frames);
    cpl_errorstate_set(es); /* swallow the illegal input state */
    cpl_msg_debug(__func__, "No raman lines found in input frameset!");
    return NULL;
  }
  const char *fn = cpl_frame_get_filename(frame);
  cpl_table *lines = muse_cpltable_load(fn, "LINES", muse_sky_lines_lines_def);
  if (!lines) {
    cpl_msg_warning(__func__, "Could not load raman lines from \"%s\"", fn);
    cpl_frameset_delete(frames);
    return NULL;
  }

  /* Output message, no matter if unit conversions happen           *
   * (so that the following messages can be related to a file name) */
  cpl_msg_info(__func__, "Loaded raman lines from \"%s\"", fn);
  /* Convert tables with unscaled flux into ones with scaled flux. */
  if (lines) {
    cpl_msg_indent_more();
    const char *unit = cpl_table_get_column_unit(lines, "flux");
    if (!unit) {
      cpl_msg_warning(__func__, "No flux unit given!");
    } else if (strcmp(unit, "erg/(s cm^2 arcsec^2)") == 0) {
      cpl_msg_info(__func__, "Scaling flux by 1e20.");
      cpl_table_multiply_scalar(lines, "flux", 1e20);
      cpl_table_set_column_unit(lines, "flux",
                                "10**(-20)*erg/(s cm^2 arcsec^2)");
    } else if (strcmp(unit, "10**(-20)*erg/(s cm^2 arcsec^2)") != 0) {
      cpl_msg_warning(__func__, "Unsupported flux unit \"%s\".", unit);
    }
    cpl_msg_indent_less();
  } /* if lines */

  muse_processing_append_used(aProcessing, frame, CPL_FRAME_GROUP_CALIB, 1);
  cpl_frameset_delete(frames);

  return lines;
} /* muse_raman_lines_load() */

/*----------------------------------------------------------------------------*/
/**
   @brief  Add a Raman line profile to pixtable
   @param  aPixtable   table with wavelengths and an lsf column
   @param  aRamanLines Raman line list
   @param  aGroup      line group to take from the Raman line list
   @param  aFlux       total flux of the line list, relative to the line list
   @param  aLsfCube    LSF data cube

   This will, slice by slice, add the raman lines profile to all
   entries in the pixtable. The pixtable is required to hold an "lsf"
   column with values set to zero. Subsequent calls of this function
   add the line spread functions of each call.
*/
/*----------------------------------------------------------------------------*/
cpl_error_code
muse_raman_add_lsf(muse_pixtable *aPixtable, cpl_table *aRamanLines,
                   int aGroup, double aFlux, muse_lsf_cube **aLsfCube)
{
  muse_pixtable **slice_pixtable = muse_pixtable_extracted_get_slices(aPixtable);

  cpl_size n_slices = muse_pixtable_extracted_get_size(slice_pixtable);
  cpl_size i_slice;
  cpl_error_code rc[n_slices];
  #pragma omp parallel for default(none)                                       \
          shared(aLsfCube, aRamanLines, aGroup, aFlux, n_slices,               \
                 slice_pixtable, rc)
  for (i_slice = 0; i_slice < n_slices; i_slice++) {
    muse_pixtable *slice_pt = slice_pixtable[i_slice];
      uint32_t origin = cpl_table_get_int(slice_pt->table, MUSE_PIXTABLE_ORIGIN,
                                          0, NULL);
    int ifu = muse_pixtable_origin_get_ifu(origin);
    int slice = muse_pixtable_origin_get_slice(origin);
    if (!aLsfCube[ifu-1]) {
      cpl_msg_warning(__func__, "No LSF params for slice #%i.%i. Ignoring in "
                      "Raman simulation.", ifu, slice);
      rc[i_slice] = CPL_ERROR_NONE;
      continue;
    }

    cpl_image *lsfImage = cpl_imagelist_get(aLsfCube[ifu-1]->img, slice-1);
    int i_line;
    int n_lines = cpl_table_get_nrow(aRamanLines);
    rc[i_slice] = CPL_ERROR_NONE;
    for (i_line = 0; i_line < n_lines; i_line++) {
      if (cpl_table_get(aRamanLines, "group", i_line, NULL) != aGroup) {
        continue;
      }
      double wavelength = cpl_table_get(aRamanLines, "lambda", i_line, NULL);
      double flux = cpl_table_get(aRamanLines, "flux", i_line, NULL) * aFlux;
      cpl_error_code res = muse_lsf_apply_slice(slice_pt, wavelength, flux,
                                                lsfImage, aLsfCube[ifu-1]->wcs);
      if (res != CPL_ERROR_NONE) {
        rc[i_slice] = res;
        break;
      }
    }
  }
  muse_pixtable_extracted_delete(slice_pixtable);

  for (i_slice = 0; i_slice < n_slices; i_slice++) {
    if (rc[i_slice] != CPL_ERROR_NONE) {
      return rc[i_slice];
    }
  }
  return CPL_ERROR_NONE;
} /* muse_raman_add_lsf() */

/*----------------------------------------------------------------------------*/
/**
   @brief  Apply the LSF and spatial Raman light distribution
   @param  aPixtable table with x and y positions, wavelength and lsf entries
   @param  aPar      Parameter table
   @param  aRetval   Array to hold the data values, or NULL.
   @return Array with fit values

   The fit uses a common spatial profile (two dimensional 2nd order
   polynomial) for both pixtables. The spatial profile is returned in
   the first five values of the returned array:

    * 0 - x² coefficient
    * 1 - xy coefficient
    * 2 - y² coefficient
    * 3 - x  coefficient
    * 4 - y  coefficient

   The flux is normalized to be 1 at the center (x=y=0).

   This requires the pixtables to have an additional "lsf" column
   containing the Line profile of the Raman scattering, normalized to
   one.
 */
/*----------------------------------------------------------------------------*/
cpl_array *
muse_raman_simulate(muse_pixtable *aPixtable, cpl_array *aPar,
                    cpl_array *aRetval)
{
  cpl_array *xpos = muse_cpltable_extract_column(aPixtable->table,
                                                 MUSE_PIXTABLE_XPOS);
  cpl_array *ypos = muse_cpltable_extract_column(aPixtable->table,
                                                 MUSE_PIXTABLE_YPOS);
  cpl_array *lsf = muse_cpltable_extract_column(aPixtable->table, "lsf");

  /* Apply parameterization:
     p[0]*xpos^2 + p[1]*xpos*ypos + p[2]*ypos^2 + p[3]*xpos + p[4]*ypos + 1
  */
  if (!aRetval) {
    aRetval = cpl_array_new(cpl_array_get_size(lsf), CPL_TYPE_DOUBLE);
  }
  cpl_array_fill_window(aRetval, 0, cpl_array_get_size(aRetval), 0);
  cpl_array *tmp;

  tmp = cpl_array_duplicate(xpos);
  cpl_array_multiply(tmp, xpos);
  cpl_array_multiply_scalar(tmp, cpl_array_get(aPar, 0, NULL));
  cpl_array_add(aRetval, tmp);
  cpl_array_delete(tmp);

  tmp = cpl_array_duplicate(xpos);
  cpl_array_multiply(tmp, ypos);
  cpl_array_multiply_scalar(tmp, cpl_array_get(aPar, 1, NULL));
  cpl_array_add(aRetval, tmp);
  cpl_array_delete(tmp);

  tmp = cpl_array_duplicate(ypos);
  cpl_array_multiply(tmp, ypos);
  cpl_array_multiply_scalar(tmp, cpl_array_get(aPar, 2, NULL));
  cpl_array_add(aRetval, tmp);
  cpl_array_delete(tmp);

  tmp = cpl_array_duplicate(xpos);
  cpl_array_multiply_scalar(tmp, cpl_array_get(aPar, 3, NULL));
  cpl_array_add(aRetval, tmp);
  cpl_array_delete(tmp);

  tmp = cpl_array_duplicate(ypos);
  cpl_array_multiply_scalar(tmp, cpl_array_get(aPar, 4, NULL));
  cpl_array_add(aRetval, tmp);
  cpl_array_delete(tmp);

  cpl_array_add_scalar(aRetval, 1);

  /* multiply with LSF */
  cpl_array_multiply(aRetval, lsf);

  cpl_array_unwrap(xpos);
  cpl_array_unwrap(ypos);
  cpl_array_unwrap(lsf);

  return aRetval;
} /* muse_raman_simulate() */

/*----------------------------------------------------------------------------*/
/**
  @brief  Simulate an image of the parameters of the spatial Raman light fit.
  @param  aImage   input image
  @param  aPar     parameter array
  @return A new muse_image * with the simulated Raman distribution.

  aImage is used to define the (pixel-table) WCS (which was used to fit the
  Raman light distribution) and as basis for the output image.

  aPar contains the fit parameters, see @ref muse_raman_fit().

  @error{set CPL_ERROR_NULL_INPUT\, return NULL, aImage and/or aPar are NULL}
  @error{set CPL_ERROR_ILLEGAL_INPUT\, return NULL,
         aPar does not contain 5 entries}
 */
/*----------------------------------------------------------------------------*/
muse_image *
muse_raman_simulate_image(const muse_image *aImage, const cpl_array *aPar)
{
  cpl_ensure(aImage && aPar, CPL_ERROR_NULL_INPUT, NULL);
  cpl_ensure(cpl_array_get_size(aPar) == 5 + MUSE_RAMAN_RANGES,
             CPL_ERROR_ILLEGAL_INPUT, NULL);

  /* get the WCS */
  muse_wcs *wcs = muse_wcs_new(aImage->header);
  if (!wcs) {
    /* not a critical error for the actual processing, so just warn */
    cpl_msg_warning(__func__, "No Raman image simulated: %s",
                    cpl_error_get_message());
    /* but still set the error code */
    cpl_error_set(__func__, cpl_error_get_code());
    return NULL;
  }
  /* create output image, just header plus data are enough here */
  muse_image *image = muse_image_new();
  image->header = cpl_propertylist_duplicate(aImage->header);
  image->data = cpl_image_duplicate(aImage->data);
  /* no bad pixels in the image we create here (they get intra/extrapolated) */
  cpl_image_accept_all(image->data);

  /* get buffers for parmeters and image */
  const double *p = cpl_array_get_data_double_const(aPar);
  float *data = cpl_image_get_data_float(image->data);
  int i, nx = cpl_image_get_size_x(image->data),
      ny = cpl_image_get_size_y(image->data);
  for (i = 0; i < nx; i++) {
    int j;
    for (j = 0; j < ny; j++) {
      /* get location in WCS coordinates */
      double x, y;
      muse_wcs_projplane_from_pixel_fast(wcs, i+1, j+1, &x, &y);

      /* evaluate the polynomial */
      data[i + j*nx] = p[0] * x*x + p[1] * x*y + p[2] * y*y + p[3] * x
                     + p[4] * y + 1.;
    } /* for j */
  } /* for i */
  cpl_free(wcs);

  return image;
} /* muse_raman_simulate_image() */

typedef struct {
  muse_pixtable *pt[MUSE_RAMAN_RANGES];
} muse_raman_fit_struct;

/*----------------------------------------------------------------------------*/
/**
   @brief Evaluate a given parameter set for the Raman fit
   @param aData     Data forwarded from the fit algorithm call (structure
                    containing the pixtables)
   @param aPar      Current fit parameter.
   @param aRetval   Return value vector.
   @return CPL_ERROR_NONE if everything went OK.
 */
/*----------------------------------------------------------------------------*/
static cpl_error_code
muse_raman_eval(void *aData, cpl_array *aPar, cpl_array* aRetval)
{
  muse_raman_fit_struct *fitData = aData;
  int i;
  cpl_size i_row = 0;
  for (i = 0; i < MUSE_RAMAN_RANGES; i++) {
    if (!fitData->pt[i]) {
      continue;
    }
    cpl_array *a =
      cpl_array_wrap_double(cpl_array_get_data_double(aRetval) + i_row,
                            cpl_table_get_nrow(fitData->pt[i]->table));

    muse_raman_simulate(fitData->pt[i], aPar, a);
    cpl_array_multiply_scalar(a, cpl_array_get(aPar, 5 + i, NULL));
    cpl_array *data = muse_cpltable_extract_column(fitData->pt[i]->table,
                                                   MUSE_PIXTABLE_DATA);
    cpl_array_subtract(a, data);
    cpl_array_unwrap(data);
    cpl_array_subtract_scalar(a, cpl_array_get_mean(a));
    cpl_array_unwrap(a);
    i_row += cpl_table_get_nrow(fitData->pt[i]->table);
  } /* for i (both Raman ranges) */

  return CPL_ERROR_NONE;
} /* muse_raman_eval() */

/*----------------------------------------------------------------------------*/
/**
   @brief  Fit raman lines to an array of pixtables
   @param  aPixtable array of MUSE_RAMAN_RANGES (2) pixtables
   @return Array with fit values

   The fit uses a common spatial profile (two dimensional 2nd order
   polynomial) for both pixtables. The spatial profile is returned in
   the first five values of the returned array:

    * 0 - x² coefficient
    * 1 - xy coefficient
    * 2 - y² coefficient
    * 3 - x  coefficient
    * 4 - y  coefficient

   The remaining two values of the array contain the central (x=y=0)
   flux for both pixtables.

   This requires the pixtables to have an additional "lsf" column
   containing the Line profile of the Raman scattering, normalized to
   one.
 */
/*----------------------------------------------------------------------------*/
cpl_array *
muse_raman_fit(muse_pixtable *aPixtable[MUSE_RAMAN_RANGES])
{
  muse_raman_fit_struct fitData;
  int i;
  for (i = 0; i < MUSE_RAMAN_RANGES; i++) {
    fitData.pt[i] = aPixtable[i];
  }

  cpl_array *param = cpl_array_new(5 + MUSE_RAMAN_RANGES, CPL_TYPE_DOUBLE);
  /* First quess. See muse_raman_polynomial() */
  cpl_array_set(param, 0, 0.); /* x² coefficient */
  cpl_array_set(param, 1, 0.); /* xy coefficient */
  cpl_array_set(param, 2, 0.); /* y² coefficient */
  cpl_array_set(param, 3, 0.); /* x  coefficient */
  cpl_array_set(param, 4, 0.); /* y  coefficient */

  for (i = 0; i < MUSE_RAMAN_RANGES; i++) {
    cpl_array_set(param, 5 + i, 0.); /* constant part i-th range*/
  }

  /* do the fit, ignoring possible errors */
  int debug = getenv("MUSE_DEBUG_LSF_FIT")
            && atoi(getenv("MUSE_DEBUG_LSF_FIT")) > 0;
  muse_cpl_optimize_control_t ctrl = {
    -1, -1, -1, -1, /* default ftol, xtol, gtol, maxiter */
    debug
  };

  cpl_size n_row = 0;
  for (i = 0; i < MUSE_RAMAN_RANGES; i++) {
    if (!aPixtable[i]) {
      continue;
    }
    n_row += cpl_table_get_nrow(aPixtable[i]->table);
  } /* for i (both Raman ranges) */
  cpl_msg_info(__func__, "Starting Raman line fit");
  cpl_error_code r = muse_cpl_optimize_lvmq(&fitData, param, n_row,
                                            muse_raman_eval, &ctrl);
  if (r != CPL_ERROR_NONE) {
    cpl_msg_warning(__func__, "Raman line fit failed: %s",
                    cpl_error_get_message());
  }

  return param;
} /* muse_raman_fit() */

/**@}*/
