/* -*- 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) 2008-2014 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.
 */

#ifdef HAVE_CONFIG_H
#include <config.h>
#endif

#include <math.h>
#include <string.h>

#include "muse_sky.h"
#include "muse_instrument.h"
#include "muse_lsf.h"
#include "muse_optimize.h"

/** @addtogroup muse_skysub
    @{
  */

/*----------------------------------------------------------------------------*/
/**
   @brief Create a new fit parameter structure
   @param aOffset Set to CPL_True if the offset is to be fit
   @param aRefraction set to CPL_True if the atmospheric refraction is to be fit
   @param aSensitivity Order of the sensitivity fit parameter
   @param aSlitWidth Order of the slit width fit parameter
   @param aBinWidth Order of the bin width fit parameter
   @param aLSFWidth Order of the LSF width fit parameter
   @param aHermit3 Order of the 3rd order hermitean fit parameter
   @param aHermit4 Order of the 4rd order hermitean fit parameter
   @param aHermit5 Order of the 5rd order hermitean fit parameter
   @param aHermit6 Order of the 6rd order hermitean fit parameter

   @return Pointer to a newly allocates fit parameter structure.
 */

/*----------------------------------------------------------------------------*/

muse_sky_fit_params *
muse_sky_fit_params_new(cpl_boolean aOffset, cpl_boolean aRefraction,
                        cpl_size aSensitivity, cpl_size aSlitWidth,
                        cpl_size aBinWidth, cpl_size aLSFWidth,
                        cpl_size aHermit3, cpl_size aHermit4,
                        cpl_size aHermit5, cpl_size aHermit6)
{
  muse_sky_fit_params *params = cpl_malloc(sizeof(muse_sky_fit_params));
  params->offset = aOffset;
  params->refraction = aRefraction;
  params->sensitivity = aSensitivity;
  params->slit_width = aSlitWidth;
  params->bin_width = aBinWidth;
  params->lsf_width = aLSFWidth;
  params->hermit[0] = aHermit3;
  params->hermit[1] = aHermit4;
  params->hermit[2] = aHermit5;
  params->hermit[3] = aHermit6;

  params->n_param = params->offset + params->refraction + params->sensitivity +
    params->slit_width + params->bin_width + params->lsf_width;
  cpl_size i;
  for (i = 0; i < MAX_HERMIT_ORDER; i++) {
    params->n_param += params->hermit[i];
  }

  return params;
}

/*----------------------------------------------------------------------------*/
/**
   @brief Delete the fit parameter structure.
   @param params Structure to destroy
*/
/*----------------------------------------------------------------------------*/
void
muse_sky_fit_params_delete(muse_sky_fit_params *params) {
  cpl_free(params);
}

/*----------------------------------------------------------------------------*/
/**
   @private
   @brief Convert a parametrization to LSF parameters.
   @param aTemplate    Template for parametrization values
   @param aPar         Parametrization values
   @param aFitParams   Structure defining slice parameter fitting flags
   @return Pointer to a newly allocates LSF parameter structure.

   This is the parametrization used in the SLICE procedure.
 */
/*----------------------------------------------------------------------------*/
muse_lsf_params *
muse_sky_slice_apply_lsf_parametrization(const muse_lsf_params *aTemplate,
                                         const cpl_array *aPar,
                                         const muse_sky_fit_params *aFitParams)
{
  cpl_size offset = 0;
  cpl_size j;

  cpl_size max_hermit = 0;
  for (j = 0; j < MAX_HERMIT_ORDER; j++) {
    if (max_hermit < aFitParams->hermit[j]) {
      max_hermit = aFitParams->hermit[j];
    }
  }
  muse_lsf_params *lsf
    = muse_lsf_params_new((aFitParams->sensitivity > 0)?
                          aFitParams->sensitivity:
                          cpl_array_get_size(aTemplate->sensitivity),
                          (aFitParams->lsf_width > 0)?
                          aFitParams->lsf_width:
                          cpl_array_get_size(aTemplate->lsf_width),
                          (max_hermit > 0)?
                          max_hermit:
                          cpl_array_get_size(aTemplate->hermit[0]));
  cpl_array_set(lsf->sensitivity, 0, 1.0);

  if (aFitParams->offset > 0) {
    lsf->offset = cpl_array_get(aPar, offset++, NULL);
  } else {
    lsf->offset = aTemplate->offset;
  }

  if (aFitParams->refraction > 0) {
    lsf->refraction = 1.0 + cpl_array_get(aPar, offset++, NULL);
  } else {
    lsf->refraction = aTemplate->refraction;
  }

  cpl_size n = cpl_array_get_size(lsf->sensitivity);
  if (aFitParams->sensitivity > 0) {
    for (j = 0; j < n; j++) {
      if (j < aFitParams->sensitivity) {
        cpl_array_set(lsf->sensitivity, j,
                      cpl_array_get(aPar, offset++, NULL));
      } else {
      cpl_array_set(lsf->sensitivity, j, 0.0);
      }
    }
  } else {
    for (j = 0; j < n; j++) {
      cpl_array_set(lsf->sensitivity, j,
                    cpl_array_get(aTemplate->sensitivity, j, NULL));
    }
  }

  if (aFitParams->slit_width > 0) {
    lsf->slit_width = cpl_array_get(aPar, offset++, NULL);
  } else {
    lsf->slit_width = aTemplate->slit_width;
  }

  if (aFitParams->bin_width > 0) {
    lsf->bin_width = cpl_array_get(aPar, offset++, NULL);
  } else {
    lsf->bin_width = aTemplate->bin_width;
  }

  n = cpl_array_get_size(lsf->lsf_width);
  if (aFitParams->lsf_width > 0) {
    for (j = 0; j < n; j++) {
      if (j < aFitParams->lsf_width) {
        cpl_array_set(lsf->lsf_width, j,
                      cpl_array_get(aPar, offset++, NULL));
      } else {
        cpl_array_set(lsf->lsf_width, j, 0.0);
      }
    }
  } else {
    for (j = 0; j < n; j++) {
      cpl_array_set(lsf->lsf_width, j,
                    cpl_array_get(aTemplate->lsf_width, j, NULL));
    }
  }

  int i;
  for (i = 0; i < MAX_HERMIT_ORDER; i++) {
    n = cpl_array_get_size(lsf->hermit[i]);
    if (aFitParams->hermit[i] > 0) {
      for (j = 0; j < n; j++) {
        if (j < aFitParams->hermit[i]) {
          cpl_array_set(lsf->hermit[i], j,
                        cpl_array_get(aPar, offset++, NULL));
        } else {
          cpl_array_set(lsf->hermit[i], j, 0.0);
        }
      }
    } else {
      for (j = 0; j < n; j++) {
        cpl_array_set(lsf->hermit[i], j,
                      //cpl_array_get(aTemplate->hermit[i], j, NULL));
                      0.0);
      }
    }
  }

  if (offset > cpl_array_get_size(aPar)) {
    cpl_msg_error(__func__,
                  "inconsistent array: size %ld, read with %ld values",
                  (long)cpl_array_get_size(aPar), (long)offset);
    muse_lsf_params_delete_one(lsf);
    return NULL;
  }

  return lsf;
}

/*----------------------------------------------------------------------------*/
/**
   @private
   @brief Return the parametrization values for the specific LSF params
   @param aLsf    slice LSF params.
   @param aFitParams   Structure defining slice parameter fitting flags
   @return An array containing the first guess values.
 */
/*----------------------------------------------------------------------------*/
cpl_array *
muse_sky_slice_lsf_set_param(muse_lsf_params *aLsf,
                             const muse_sky_fit_params *aFitParams) {
  cpl_array *pars = cpl_array_new(aFitParams->n_param, CPL_TYPE_DOUBLE);
  cpl_size offset = 0;

  // Relative refraction ratio - 1
  if (aFitParams->offset > 0) {
    cpl_array_set(pars, offset++, aLsf->offset);
  }
  if (aFitParams->refraction > 0) {
    cpl_array_set(pars, offset++, aLsf->refraction -1);
  }
  // relative sensitivity
  cpl_size j;
  cpl_size n = cpl_array_get_size(aLsf->sensitivity);
  for (j = 0; j < aFitParams->sensitivity; j++) {
    if (j < n) {
      cpl_msg_debug(__func__, "S[%li]=%f", (long)j,
                    cpl_array_get(aLsf->sensitivity, j, NULL));
      cpl_array_set(pars, offset++,
                    cpl_array_get(aLsf->sensitivity, j, NULL));
    } else {
      cpl_array_set(pars, offset++, (j == 0)?1.0:0.0);
    }
  }

  if (aFitParams->slit_width > 0) {
    cpl_array_set(pars, offset++, aLsf->slit_width);
  }

  if (aFitParams->bin_width > 0) {
    cpl_array_set(pars, offset++, aLsf->bin_width);
  }

  // LSF width
  n = cpl_array_get_size(aLsf->lsf_width);
  for (j = 0; j < aFitParams->lsf_width; j++) {
    if (j < n) {
      cpl_array_set(pars, offset++,
                    cpl_array_get(aLsf->lsf_width, j, NULL));
    } else {
      cpl_array_set(pars, offset++, (j == 0)?1.0:0.0);
    }
  }

  // Hermitean coefficients
  cpl_size i;
  for (i = 0; i < MAX_HERMIT_ORDER; i++) {
    n = cpl_array_get_size(aLsf->hermit[i]);
    for (j = 0; j < aFitParams->hermit[i]; j++) {
      if (j < n) {
        cpl_array_set(pars, offset++,
                      cpl_array_get(aLsf->hermit[i], j, NULL));
      } else {
        cpl_array_set(pars, offset++, 0.0);
      }
    }
  }

  if (offset > cpl_array_get_size(pars)) {
    cpl_msg_error(__func__,
                  "inconsistent array: size %ld, filled with %ld values",
                  (long)cpl_array_get_size(pars), (long)offset);
  }
  return pars;
}

/*----------------------------------------------------------------------------*/
/**
   @private
   @brief Return the parametrization values of the first guess.
   @param aFitParams    Structure defining slice parameter fitting flags
   @return An array containing the first guess values.

   This array can be converted to lsf parametern using
   muse_lsf_slice_apply_lsf_parametrization().
 */
/*----------------------------------------------------------------------------*/
cpl_array *
muse_sky_slice_lsf_firstguess(const muse_sky_fit_params *aFitParams) {
  cpl_array *pars = cpl_array_new(aFitParams->n_param, CPL_TYPE_DOUBLE);
  cpl_size offset = 0;

  // Wavelength offset
  if (aFitParams->offset > 0) {
    cpl_array_set(pars, offset++, 0.0);
  }

  // Relative refraction ratio - 1
  if (aFitParams->refraction > 0) {
    cpl_array_set(pars, offset++, 0.0);
  }

  // Relative sensitivity
  cpl_size j;
  for (j = 0; j < aFitParams->sensitivity; j++) {
    cpl_array_set(pars, offset++, (j == 0)?1.0:0.0);
  }

  // Slit width
  if (aFitParams->slit_width > 0) {
    cpl_array_set(pars, offset++, kMuseSliceSlitWidthA);
  }

  // Bin width
  if (aFitParams->bin_width > 0) {
    cpl_array_set(pars, offset++, kMuseSpectralSamplingA);
  }

  // LSF width
  for (j = 0; j < aFitParams->lsf_width; j++) {
    cpl_array_set(pars, offset++, (j == 0)?0.5:0.0);
  }

  // Hermitean coefficients
  cpl_size i;
  for (i = 0; i < MAX_HERMIT_ORDER; i++) {
    for (j = 0; j < aFitParams->hermit[i]; j++) {
      cpl_array_set(pars, offset++, 0.0);
    }
  }

  if (offset > cpl_array_get_size(pars)) {
    cpl_msg_error(__func__,
                  "inconsistent array: size %ld, filled with %ld values",
                  (long)cpl_array_get_size(pars), (long)offset);
  }
  return pars;
}

typedef struct {
  cpl_array *lambda;
  cpl_array *values;
  cpl_array *stat;
  cpl_array *sensitivity;
  cpl_array *continuum;
  cpl_table *lines;
  const muse_sky_fit_params *fit_params;
  muse_lsf_params *firstGuess;
} muse_slice_fit_struct;

/*----------------------------------------------------------------------------*/
/**
  @private
   @brief Evaluate a given parameter set
   @param aData Data forwarded from the fit algorithm call.
   @param aPar  Current fit parameter.
   @param aRetval Return value vector.
   @return CPL_ERROR_NONE if everything went OK.
 */
/*----------------------------------------------------------------------------*/
static cpl_error_code
muse_sky_slice_eval(void *aData, cpl_array *aPar, cpl_array *aRetval) {

  muse_slice_fit_struct *data = aData;
  cpl_size size = cpl_array_get_size(aRetval);

  muse_lsf_params *lsf
    = muse_sky_slice_apply_lsf_parametrization(data->firstGuess,
                                               aPar, data->fit_params);

  cpl_array *simulated = muse_sky_apply_lsf(data->lambda, data->lines, lsf);

  muse_lsf_params_delete_one(lsf);

  cpl_array_subtract(simulated, data->values);
  cpl_array_fill_window_double(aRetval, 0, size, 0.0);
  memcpy(cpl_array_get_data_double(aRetval),
         cpl_array_get_data_double_const(simulated),
         size * sizeof(double));

  cpl_array_delete(simulated);

  return CPL_ERROR_NONE;
}

/*----------------------------------------------------------------------------*/
/**
   @brief Fit all entries of one slice.
   @param aLambda Wavelength array [Angstrom].
   @param aData Measured spectrum.
   @param aStat The variance spectrum.
   @param aLines @ref muse_sky_lines_lines_def "List of emission lines".
   @param aContinuum Continuum spectrum, or NULL
   @param aSensitivity Sensitivity array, must correspond to aLambda.
   @param aFirstGuess First guess for LSF params, or NULL.
   @param aMaxIter Maximum number of iterations.
   @param aFitParams Specification which parameters to fit
   @return The fitted LSF parameters
 */
/*----------------------------------------------------------------------------*/
muse_lsf_params *
muse_sky_lsf_params_fit(const cpl_array *aLambda, const cpl_array *aData,
                        const cpl_array *aStat, const cpl_table *aLines,
                        const cpl_array *aContinuum,
                        const cpl_array *aSensitivity,
                        muse_lsf_params *aFirstGuess, int aMaxIter,
                        const muse_sky_fit_params *aFitParams) {

  muse_lsf_params *firstGuess = (aFirstGuess != NULL)?
    aFirstGuess: muse_lsf_params_new(1, 3, 1);
  cpl_size reduction = 5; // we use only 20% of all pixels for the fit (speed).
  cpl_size size = cpl_array_get_size(aLambda)/reduction;
  cpl_array *lambda = cpl_array_new(size, CPL_TYPE_DOUBLE);
  cpl_array *data = cpl_array_new(size, CPL_TYPE_DOUBLE);
  cpl_array *continuum = NULL;
  if (aContinuum != NULL) {
    continuum = cpl_array_new(size, CPL_TYPE_DOUBLE);
  }
  cpl_array *sensitivity = NULL;
  if (aSensitivity != NULL) {
    sensitivity = cpl_array_new(size, CPL_TYPE_DOUBLE);
  }
  cpl_size i;
  for (i = 0; i < size; i++) {
    cpl_array_set(lambda, i, cpl_array_get(aLambda, reduction * i, NULL));
    cpl_array_set(data, i, cpl_array_get(aData, reduction * i, NULL));
    if (aContinuum != NULL) {
      cpl_array_set(continuum, i, cpl_array_get(aContinuum, reduction * i, NULL));
    }
    if (aSensitivity != NULL) {
      cpl_array_set(sensitivity, i,
                    cpl_array_get(aSensitivity, reduction * i, NULL));
    }
  }

  cpl_array *stat = cpl_array_new(size, CPL_TYPE_DOUBLE);
  for (i = 0; i < size; i++) {
    cpl_array_set(stat, i, sqrt(cpl_array_get(aStat, reduction * i, NULL)));
  }

  cpl_table *lines = cpl_table_duplicate(aLines);
  double maxflux = cpl_table_get_column_max(lines, "flux");
  muse_sky_lines_cut(lines, 3e-3 * maxflux);

  muse_slice_fit_struct fit_data = {
    lambda,
    data,
    stat,
    sensitivity,
    continuum,
    lines,
    aFitParams,
    firstGuess,
  };

  cpl_array *pars = (aFirstGuess == NULL)?
    muse_sky_slice_lsf_firstguess(aFitParams):
    muse_sky_slice_lsf_set_param(aFirstGuess, aFitParams);

  cpl_error_code r = CPL_ERROR_NONE;
  int debug = getenv("MUSE_DEBUG_LSF_FIT")
            && atoi(getenv("MUSE_DEBUG_LSF_FIT")) > 0;
  if (aMaxIter > 0) {
    muse_cpl_optimize_control_t ctrl = {
      -1, -1, -1, // default ftol, xtol, gtol
      aMaxIter, debug
    };

    r = muse_cpl_optimize_lvmq(&fit_data, pars, size,
                               muse_sky_slice_eval, &ctrl);
  }

  if (r != CPL_ERROR_NONE) { // on error: reset to first guess
    cpl_array_delete(pars);
    pars = (aFirstGuess == NULL)?
      muse_sky_slice_lsf_firstguess(aFitParams):
      muse_sky_slice_lsf_set_param(aFirstGuess, aFitParams);
  }

  muse_lsf_params *lsf
    = muse_sky_slice_apply_lsf_parametrization(firstGuess, pars, aFitParams);

  cpl_array_delete(pars);
  cpl_table_delete(fit_data.lines);
  cpl_array_delete(fit_data.lambda);
  cpl_array_delete(fit_data.sensitivity);
  cpl_array_delete(fit_data.continuum);
  cpl_array_delete(fit_data.values);
  cpl_array_delete(fit_data.stat);
  if (aFirstGuess == NULL) {
    // remove temporary first guess template
    muse_lsf_params_delete_one(firstGuess);
  }

  return lsf;
}

/**@}*/
