/* -*- 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"
#include "muse_utils.h"
#include "muse_data_format_z.h"

/** @addtogroup muse_skysub */
/**@{*/

/*----------------------------------------------------------------------------*/
/**
   @private
   @brief Return the parametrization values of the first guess.
   @param ngroups   Number of line groups.
   @return An array containing the first guess values.

   This array can be converted to sky parametern using
   muse_sky_master_apply_parametrization().
 */
/*----------------------------------------------------------------------------*/
static cpl_array *
muse_sky_lines_firstguess(int ngroups) {
  cpl_array *pars = cpl_array_new(ngroups+2, CPL_TYPE_DOUBLE);

  // Line strengths
  cpl_size i;
  for (i = 0; i < ngroups; i++) {
    cpl_array_set(pars, i, 1e-1);
  }
  cpl_array_set(pars, ngroups, 0.0);
  cpl_array_set(pars, ngroups+1, 0.0);
  return pars;
}

typedef struct {
  const cpl_array *lambda; // wavelength
  const cpl_array *values; // data values
  const cpl_array *weights; // data weights
  const cpl_table *lines; // constant sky physics data
  const cpl_size ngroups; // number of groups in the data
  const cpl_image *lsfImage;
  const muse_wcs *wcs;
} muse_master_fit_struct;

/*----------------------------------------------------------------------------*/
/**
   @private
   @brief  Simulate the spectrum for a given parameter set.
   @param  aFitData   Structure of all the sky-related data to fit.
   @param  aPar       Parameter array for relative line strengts.
   @return the simulated spectrum
 */
/*----------------------------------------------------------------------------*/
static cpl_array *
simulate_master_sky_parameters(muse_master_fit_struct *aFitData,
                               const cpl_array *aPar)
{
  cpl_table *lines = cpl_table_duplicate(aFitData->lines);
  int res = 0;
  double fac = cpl_array_get(aPar, cpl_array_get_size(aPar)-2, &res);
  double offset = cpl_array_get(aPar, cpl_array_get_size(aPar)-1, &res);
  cpl_table_multiply_scalar(lines, "lambda", 1+fac);
  cpl_table_add_scalar(lines, "lambda", offset-7000 * fac);
  cpl_array *line_strengths = cpl_array_duplicate(aPar);
  cpl_array_multiply(line_strengths, aPar);
  muse_sky_lines_apply_strength(lines, line_strengths);
  cpl_array_delete(line_strengths);

  double maxflux = cpl_table_get_column_max(lines, "flux");
  muse_sky_lines_cut(lines, 1e-4 * maxflux);

  cpl_array *simulated
    = muse_sky_lines_spectrum(aFitData->lambda, lines,
                              aFitData->lsfImage, aFitData->wcs);
  cpl_table_delete(lines);

  return simulated;
}

/*----------------------------------------------------------------------------*/
/**
   @private
   @brief Evaluate a given parameter set for the MASTER SKY calc.
   @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_master_eval(void *aData, cpl_array *aPar, cpl_array *aRetval)
{
  cpl_size size = cpl_array_get_size(aRetval);
  muse_master_fit_struct *fitData = aData;
  cpl_array *simulated = simulate_master_sky_parameters(fitData, aPar);
  cpl_array_subtract(simulated, fitData->values);
  cpl_array *dsimulated = muse_cplarray_diff(simulated, 1);

  cpl_array_multiply(dsimulated, fitData->weights);

  cpl_array_fill_window_double(aRetval, 0, size, 0.0);
  memcpy(cpl_array_get_data_double(aRetval),
         cpl_array_get_data_double_const(dsimulated),
         size * sizeof(double));
  cpl_array_delete(simulated);
  cpl_array_delete(dsimulated);

  return CPL_ERROR_NONE;
}

/*----------------------------------------------------------------------------*/
/**
   @private
   @brief  Correct the first guess by data.
   @param  aFitData   Structure of all the sky-related data to fit.
   @param  aPar       Parameter array for relative line strengts.
   @return CPL_ERROR_NONE

   This function computes the spectrum from aPars, takes the bin for
   the strongest line for each group and compares it with the data
   value at that bin. Their ratio gives the correction factor for this
   line group strength.
 */
/*----------------------------------------------------------------------------*/
static cpl_error_code
muse_sky_lines_correct_firstguess(muse_master_fit_struct *aFitData,
                                  cpl_array *aPar)
{
  // evaluate the first guess
  cpl_array *simulated = simulate_master_sky_parameters(aFitData, aPar);

  cpl_table *lines = cpl_table_duplicate(aFitData->lines);
  cpl_array *line_strengths = cpl_array_duplicate(aPar);
  cpl_array_multiply(line_strengths, aPar);
  muse_sky_lines_apply_strength(lines, line_strengths);
  cpl_array_delete(line_strengths);
  cpl_size i_group;
  for (i_group = 0; i_group < aFitData->ngroups; i_group++) {
    // take the strongest line of each group
    cpl_table_unselect_all(lines);
    cpl_table_or_selected_int(lines, "group", CPL_EQUAL_TO, i_group);
    cpl_table *gtable = cpl_table_extract_selected(lines);
    if (cpl_table_get_nrow(gtable) == 0) {
      cpl_table_delete(gtable);
      continue;
    }
    cpl_size row;
    cpl_table_get_column_maxpos(gtable, "flux", &row);
    double wavelength = cpl_table_get_double(gtable, "lambda", row, NULL);

    // divide measured data and first guess result
    cpl_size i_lbda1 = muse_cplarray_find_sorted(aFitData->lambda,
                                                 wavelength - 3.0);
    cpl_size i_lbda2 = muse_cplarray_find_sorted(aFitData->lambda,
                                                 wavelength + 3.0);
    double y_data = 0;
    double y_sim = 0;
    double offset = cpl_array_get(aFitData->values, i_lbda1, NULL);
    offset += cpl_array_get(aFitData->values, i_lbda2, NULL);
    offset /= 2;
    cpl_size i_lbda;
    for (i_lbda = i_lbda1; i_lbda <= i_lbda2; i_lbda++) {
      double wy_data = cpl_array_get(aFitData->values, i_lbda, NULL) - offset;
      if (wy_data < 0) wy_data = 0;
      double wy_sim = cpl_array_get(simulated, i_lbda, NULL);
      y_data += wy_data;
      y_sim += wy_sim;
    }

    // take this as correction factor
    if (y_sim > 0) {
      double y = cpl_array_get(aPar, i_group, NULL);
      cpl_array_set(aPar, i_group, y*sqrt(y_data/y_sim));
    }
    cpl_table_delete(gtable);
  }
  cpl_table_delete(lines);
  cpl_array_delete(simulated);

  return CPL_ERROR_NONE;
}

/*----------------------------------------------------------------------------*/
/**
   @brief Fit all entries of the pixel table to the master sky.
   @param aSpectrum   Sky spectrum
   @param aLines      Sky lines table. The line fluxes are adjusted.
   @param aLsfImage   Image to store the LSF.
   @param aWCS        The WCS header of the image.
   @return CPL_ERROR_NONE if everythin is OK.

   @cpl_ensure_code{aSpectrum != NULL, CPL_ERROR_NULL_INPUT}
   @cpl_ensure_code{aLines != NULL, CPL_ERROR_NULL_INPUT}
   @error{set CPL_ERROR_DATA_NOT_FOUND, aSpectrum has no entries}
 */
/*----------------------------------------------------------------------------*/
cpl_error_code
muse_sky_lines_fit(cpl_table *aSpectrum, cpl_table *aLines,
                   cpl_image *aLsfImage, muse_wcs *aWCS)
{
  cpl_ensure_code(aSpectrum, CPL_ERROR_NULL_INPUT);
  cpl_ensure_code(aLines, CPL_ERROR_NULL_INPUT);
  cpl_size nRows = cpl_table_get_nrow(aSpectrum);
  cpl_ensure_code(nRows > 0, CPL_ERROR_DATA_NOT_FOUND);

  cpl_array *lambda = muse_cpltable_extract_column(aSpectrum, "lambda");
  cpl_array *data = muse_cpltable_extract_column(aSpectrum, "data");
  cpl_array *stat = muse_cpltable_extract_column(aSpectrum, "stat");

  cpl_array *weights = cpl_array_extract(stat, 0, nRows - 1);
  cpl_array *w2 = cpl_array_extract(stat, 1, nRows);
  cpl_array_add(weights, w2);
  cpl_array_delete(w2);
  cpl_array_power(weights, -0.5);
  //  cpl_array_fill_window(weights, 0, cpl_array_get_size(weights)-1, 1.0);

  muse_master_fit_struct fit_data = {
    lambda,
    data,
    weights,
    aLines,
    cpl_table_get_column_max(aLines, "group") + 1,
    aLsfImage,
    aWCS
  };

  cpl_array *pars = muse_sky_lines_firstguess(fit_data.ngroups);
  muse_sky_lines_correct_firstguess(&fit_data, pars);

  // 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
  };
  /* this potentially takes a long time, better output something */
  cpl_msg_info(__func__, "Starting sky line fit");
  cpl_error_code r = muse_cpl_optimize_lvmq(&fit_data, pars, nRows-1,
                                            muse_sky_master_eval, &ctrl);
  if (r != CPL_ERROR_NONE) {
    cpl_msg_error(__func__, "Sky line fit failed with error code %i: %s",
                  r, cpl_error_get_message());
  } else {
    int res = 0;
    double fac = cpl_array_get(pars, cpl_array_get_size(pars)-2, &res);
    double offset = cpl_array_get(pars, cpl_array_get_size(pars)-1, &res);
    cpl_table_multiply_scalar(aLines, "lambda", 1+fac);
    cpl_table_add_scalar(aLines, "lambda", offset-7000 * fac);
    double l_min = cpl_table_get_column_min(aLines, "lambda");
    double l_max = cpl_table_get_column_max(aLines, "lambda");
    cpl_msg_info(__func__, "Sky line fit finished successfully. "
                 "Offset %.3f A (at %.0f A) ... %.3f A (at %.0f A)",
                 offset + (l_min - 7000) *fac, l_min,
                 offset + (l_max - 7000) * fac, l_max);
  }

  cpl_array_multiply(pars, pars);
  muse_sky_lines_apply_strength(aLines, pars);
  cpl_array_delete(pars);
  // sort by flux
  cpl_propertylist *order = cpl_propertylist_new();
  cpl_propertylist_append_bool(order, "flux", TRUE);
  cpl_table_sort(aLines, order);
  cpl_propertylist_delete(order);
  cpl_array_unwrap(lambda);
  cpl_array_unwrap(data);
  cpl_array_unwrap(stat);
  cpl_array_delete(weights);

  return CPL_ERROR_NONE;
}

/*----------------------------------------------------------------------------*/
/**
   @brief Create a continuum spectrum.
   @param aSpectrum   Sky spectrum
   @param aLines      Sky emission line list
   @param aLsfImage   LSF data to use
   @param aLsfWCS     WCS information for the LSF data
   @param aBinWidth   Continuum spectrum bin with
   @return Continuum spectrum table.
   @cpl_ensure{aSpectrum, CPL_ERROR_NULL_INPUT, NULL}
   @cpl_ensure{aLines, CPL_ERROR_NULL_INPUT, NULL}
   @cpl_ensure{aLsfImage, CPL_ERROR_NULL_INPUT, NULL}
   @cpl_ensure{aLsfWCS, CPL_ERROR_NULL_INPUT, NULL}
 */
/*----------------------------------------------------------------------------*/
cpl_table *
muse_sky_continuum_create(cpl_table *aSpectrum, cpl_table *aLines,
                          cpl_image *aLsfImage, muse_wcs *aLsfWCS,
                          double aBinWidth)
{
  cpl_ensure(aSpectrum, CPL_ERROR_NULL_INPUT, NULL);
  cpl_ensure(aLines, CPL_ERROR_NULL_INPUT, NULL);
  cpl_ensure(aLsfImage, CPL_ERROR_NULL_INPUT, NULL);
  cpl_ensure(aLsfImage, CPL_ERROR_NULL_INPUT, NULL);

  cpl_array *lambda = muse_cpltable_extract_column(aSpectrum, "lambda");
  cpl_array *flux = muse_cpltable_extract_column(aSpectrum, "data");
  cpl_array *simulated = muse_sky_lines_spectrum(lambda, aLines, aLsfImage, aLsfWCS);
  cpl_array_subtract(simulated, flux);
  cpl_array_multiply_scalar(simulated, -1.);

  double l_min = cpl_array_get_min(lambda);
  double l_max = cpl_array_get_max(lambda);
  cpl_size n_new = (l_max - l_min) / aBinWidth;
  cpl_table *continuum = muse_cpltable_new(muse_fluxspectrum_def, n_new);
  cpl_table_fill_column_window(continuum, "flux", 0, n_new, 0.0);
  cpl_array *lambda_new = muse_cpltable_extract_column(continuum, "lambda");

  cpl_size i;
  for (i = 0; i < n_new; i++) {
    cpl_table_set(continuum, "lambda", i, l_min + i * aBinWidth);
  }
  cpl_array *flux_new = muse_cplarray_interpolate_linear(lambda_new,
                                                         lambda, simulated);
  // Copy the values to the table column
  memcpy(cpl_table_get_data_double(continuum, "flux"),
         cpl_array_get_data_double(flux_new),
         n_new * sizeof(double));
  cpl_array_delete(simulated);
  cpl_array_unwrap(lambda);
  cpl_array_unwrap(flux);
  cpl_array_unwrap(lambda_new);
  cpl_array_delete(flux_new);
  return continuum;
}

/*----------------------------------------------------------------------------*/
/**
  @brief    Load the sky continuum
  @param    aProcessing   the processing structure
  @return   The spectrum as cpl_table, or NULL
 */
/*----------------------------------------------------------------------------*/
cpl_table *
muse_sky_continuum_load(muse_processing *aProcessing)
{
  cpl_ensure(aProcessing, CPL_ERROR_NULL_INPUT, NULL);

  /* search for frames with the tag for sky continuum set */
  cpl_frameset *frames_c = muse_frameset_find(aProcessing->inframes,
                                              MUSE_TAG_SKY_CONT, 0, CPL_FALSE);
  if (frames_c == NULL || cpl_frameset_get_size(frames_c) < 1) {
    cpl_frameset_delete(frames_c);
    cpl_msg_debug(__func__, "No sky continuum found in input frameset!");
    return NULL;
  }
  cpl_frame *frame_c = cpl_frameset_get_position(frames_c, 0);
  const char *fn = cpl_frame_get_filename(frame_c);
  cpl_table *continuum = muse_cpltable_load(fn, "CONTINUUM",
                                            muse_fluxspectrum_def);

  if (continuum == NULL) {
    cpl_msg_warning(__func__, "Could not load sky continuum from \"%s\"", fn);
    cpl_frameset_delete(frames_c);
    return NULL;
  }

  cpl_msg_info(__func__, "Loaded sky continuum from \"%s\"", fn);
  muse_processing_append_used(aProcessing, frame_c, CPL_FRAME_GROUP_CALIB, 1);
  cpl_frameset_delete(frames_c);
  return continuum;
} /* muse_sky_continuum_load() */

/**@}*/
