/* -*- 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) 2005-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

/*----------------------------------------------------------------------------*
 *                             Includes                                       *
 *----------------------------------------------------------------------------*/
#include <math.h>
#include <string.h>

#include <cpl.h>

#include "muse_sky.h"
#include "muse_instrument.h"
#include "muse_quality.h"
#include "muse_tracing.h"
#include "muse_utils.h"

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

/*----------------------------------------------------------------------------*
 *                             Debugging Macros                               *
 *         Set these to 1 or higher for (lots of) debugging output            *
 *----------------------------------------------------------------------------*/
#define SKY_ROWBYROW_DEBUG 0 /* debug output in muse_sky_subtract_rowbyrow() */

/*----------------------------------------------------------------------------*/
/**
  @brief  Prepare an (object) mask for the sky row-by-row fitting.
  @param  aImage    the image to create the mask for
  @param  aTrace    the tracing table
  @return CPL_ERROR_NONE on success another cpl_error_code on failure.

  This function is based on the method developed by Lutz Wisotzki end of 2010
  and documented in Dec. 2011: detect continuum objects by vertically medianing
  Masking object pixels rejected from an iterative first-order fit in each
  slice. The object pixels are marked as EURO3D_OBJECT in the DQ extension of
  the input image.

  To find the slices on the image, this functions needs a trace table (aTrace).
  Alternatively, if aTrace is NULL, aImage needs to contain header keywords of
  the form ESO.DRS.MUSE.SLICEi.CENTER containing the approximate center of each
  slice. In that case, the row-by-row fit extends between the pixels flagged as
  EURO3D_MISSDATA on either side of this center.

  @error{return CPL_ERROR_NULL_INPUT, aImage is NULL}
  @error{return CPL_ERROR_ILLEGAL_INPUT,
         no slices are found\, because aTrace is NULL or faulty\, and aImage does not contain slice center keywords}
 */
/*----------------------------------------------------------------------------*/
cpl_error_code
muse_sky_subtract_rowbyrow_mask(muse_image *aImage, cpl_table *aTrace)
{
  cpl_ensure_code(aImage, CPL_ERROR_NULL_INPUT);
  unsigned short nslices = 0;
  if (aTrace) {
    nslices = cpl_table_get_nrow(aTrace);
  } else if (!aTrace) { /* check for the slice center headers */
    unsigned short slice = 0;
    cpl_boolean exists = CPL_FALSE;
    do {
      char *keyword = cpl_sprintf("ESO DRS MUSE SLICE%hu CENTER", ++slice);
      exists = cpl_propertylist_has(aImage->header, keyword);
#if 0
      cpl_msg_debug(__func__, "%s %s", keyword, exists ? "exists" : "doesn't exist");
#endif
      cpl_free(keyword);
    } while (exists);
    nslices = slice - 1;
  }
  cpl_msg_debug(__func__, "Found %hu slices", nslices);
  cpl_ensure_code(nslices, CPL_ERROR_ILLEGAL_INPUT);
  const unsigned int kOrdObj = 1; /* a linear fit to reject objects */
  const float kRSigObj = 3.; /* use 3 sigma rejection to find objects */

  /* vertical median as object mask */
  cpl_image *objmask = cpl_image_collapse_median_create(aImage->data, 0, 0, 0);
#if SKY_ROWBYROW_DEBUG
  cpl_msg_debug(__func__, "saving \"sky_removed_objmask.fits\"");
  cpl_image_save(objmask, "sky_removed_objmask.fits", CPL_TYPE_UNSPECIFIED,
                 NULL, CPL_IO_CREATE);
#endif
  /* get image properties and data */
  int nx = cpl_image_get_size_x(aImage->data),
      ny = cpl_image_get_size_y(aImage->data);
  int *dq = cpl_image_get_data_int(aImage->dq);

  /* loop through all slices to fit and mask the objects */
  unsigned short islice;
  for (islice = 0; islice < nslices; islice++) {
    cpl_msg_debug(__func__, "Processing slice %d", (int)islice + 1);

    /* find slice edge, either through trace info or *
     * through centers plus "missing data" flags     */
    cpl_polynomial **ptrace = NULL;
    int i1 = 0, i2 = nx, istart = 0;
    if (aTrace) {
      /* get the tracing polynomials for this slice */
      ptrace = muse_trace_table_get_polys_for_slice(aTrace, islice + 1);
      if (!ptrace) {
        cpl_msg_warning(__func__, "slice %2d: tracing polynomials missing!",
                        islice + 1);
        continue;
      }
      /* use trace of vertical center, be generous with the edges */
      i1 = floor(cpl_polynomial_eval_1d(ptrace[MUSE_TRACE_LEFT], ny/2, NULL));
      i2 = ceil(cpl_polynomial_eval_1d(ptrace[MUSE_TRACE_RIGHT], ny/2, NULL));
      if (i1 < 1 || i2 > nx || i1 > i2) {
        cpl_msg_warning(__func__, "slice %2d: faulty polynomial detected at "
                        "y=%d (borders: %d ... %d)", islice + 1, ny/2, i1, i2);
        muse_trace_polys_delete(ptrace);
        continue; /* next slice */
      }
    } else { /* find slice edge without tracing info */
      char *keyword = cpl_sprintf("ESO DRS MUSE SLICE%d CENTER", (int)islice + 1);
      istart = cpl_propertylist_get_float(aImage->header, keyword);
      cpl_free(keyword);
      /* move outwards */
      int ix = istart;
      while (!(dq[(--ix-1) + ny/2*nx] & EURO3D_MISSDATA)) {
        i1 = ix;
      }
      ix = istart;
      while (!(dq[(++ix-1) + ny/2*nx] & EURO3D_MISSDATA)) {
        i2 = ix;
      }
    } /* else, no aTrace */

    /*  fit and apply object mask */
    cpl_matrix *p = cpl_matrix_new(1, i2 - i1 + 1);
    cpl_vector *v = cpl_vector_new(i2 - i1 + 1);
    int i;
    for (i = i1; i <= i2; i++) { /* x pos start at 1 */
      cpl_matrix_set(p, 0, i - i1, i);
      int err;
      double value = cpl_image_get(objmask, i, 1, &err);
      if (err) {
        cpl_vector_set(v, i - i1, NAN); /* set to bad */
      } else {
        cpl_vector_set(v, i - i1, value);
      }
    } /* for i (relevant columns) */
    /* do iterative low-order polynomial fit (originally 1st order) */
    cpl_polynomial *p1 = muse_utils_iterate_fit_polynomial(p, v, NULL, NULL,
                                                           kOrdObj, kRSigObj,
                                                           NULL, NULL);
    int idxold = 0;
    for (i = i1; i <= i2; i++) { /* x pos start at 1 */
      int idx = idxold,
          ncol = cpl_matrix_get_ncol(p);
      while (idx < ncol && (int)cpl_matrix_get(p, 0, idx) < i) { /* object pixels */
        idx++;
      }
      /* check if we found the entry, if so skip the masking loop */
      if (idx < ncol && (int)cpl_matrix_get(p, 0, idx) == i) {
        idxold = idx; /* matrix is sorted, so keep track of position */
        continue;
      }
      int j;
      for (j = 0; j < ny; j++) { /* set all column pixels to "object" in DQ */
        dq[(i-1) + j*nx] |= EURO3D_OBJECT;
      } /* for j (vertical pixels) */
    } /* for i (relevant columns) */
    cpl_vector_delete(v);
    cpl_matrix_delete(p);
    cpl_polynomial_delete(p1);
    muse_trace_polys_delete(ptrace); /* NULL check in there... */
  } /* for islice */
  cpl_image_delete(objmask);

  return CPL_ERROR_NONE;
} /* muse_sky_subtract_rowbyrow_mask() */

/*----------------------------------------------------------------------------*/
/**
  @brief  Subtract the sky row-by-row from a CCD-based image.
  @param  aImage    the image to subtract the sky from
  @param  aTrace    the tracing table
  @param  aRSigma   rejection sigma for the iterative polynomial fit
  @param  aOrder    polynomial order for the fit of each row
  @return CPL_ERROR_NONE on success another cpl_error_code on failure.

  This function is based on the method developed by Lutz Wisotzki end of 2010
  and documented in Dec. 2011. This function expects non-relevant pixels to be
  masked in some way, i.e. only pixels with DQ values equal to EURO3D_GOODPIXEL
  are used in the fit, see @ref muse_sky_subtract_rowbyrow().
  Then go row by row through each slice on the raw image, do an iterative
  horizontal polynomial fit, and subtract it.

  To find the slices on the image, this functions needs a trace table (aTrace).
  Alternatively, if aTrace is NULL, aImage needs to contain header keywords of
  the form ESO.DRS.MUSE.SLICEi.CENTER containing the approximate center of each
  slice. In that case, the row-by-row fit extends between the pixels flagged as
  EURO3D_MISSDATA on either side of this center.

  @error{return CPL_ERROR_NULL_INPUT, aImage is NULL}
  @error{return CPL_ERROR_ILLEGAL_INPUT,
         no slices are found\, because aTrace is NULL or faulty\, and aImage does not contain slice center keywords}
 */
/*----------------------------------------------------------------------------*/
cpl_error_code
muse_sky_subtract_rowbyrow(muse_image *aImage, cpl_table *aTrace, float aRSigma,
                           unsigned int aOrder)
{
  cpl_ensure_code(aImage, CPL_ERROR_NULL_INPUT);
  unsigned short nslices = 0;
  if (aTrace) {
    nslices = cpl_table_get_nrow(aTrace);
  } else if (!aTrace) { /* check for the slice center headers */
    unsigned short slice = 0;
    cpl_boolean exists = CPL_FALSE;
    do {
      char *keyword = cpl_sprintf("ESO DRS MUSE SLICE%hu CENTER", ++slice);
      exists = cpl_propertylist_has(aImage->header, keyword);
#if 0
      cpl_msg_debug(__func__, "%s %s", keyword, exists ? "exists" : "doesn't exist");
#endif
      cpl_free(keyword);
    } while (exists);
    nslices = slice - 1;
  }
  cpl_msg_debug(__func__, "Found %hu slices", nslices);
  cpl_ensure_code(nslices, CPL_ERROR_ILLEGAL_INPUT);

  /* get image properties and data */
  int nx = cpl_image_get_size_x(aImage->data),
      ny = cpl_image_get_size_y(aImage->data);
  float *data = cpl_image_get_data_float(aImage->data),
        *stat = cpl_image_get_data_float(aImage->stat);
  int *dq = cpl_image_get_data_int(aImage->dq);

  /* loop through all slices to fit and remove the sky */
  unsigned short islice;
  for (islice = 0; islice < nslices; islice++) {
    cpl_msg_debug(__func__, "Processing slice %d", (int)islice + 1);

    /* find slice edge, either through trace info or *
     * through centers plus "missing data" flags     */
    cpl_polynomial **ptrace = NULL;
    int i1 = 0, i2 = nx, istart = 0;
    if (aTrace) {
      /* get the tracing polynomials for this slice */
      ptrace = muse_trace_table_get_polys_for_slice(aTrace, islice + 1);
      if (!ptrace) {
        cpl_msg_warning(__func__, "slice %2d: tracing polynomials missing!",
                        (int)islice + 1);
        continue;
      }
      /* use trace of vertical center, be generous with the edges */
      i1 = floor(cpl_polynomial_eval_1d(ptrace[MUSE_TRACE_LEFT], ny/2, NULL));
      i2 = ceil(cpl_polynomial_eval_1d(ptrace[MUSE_TRACE_RIGHT], ny/2, NULL));
      if (i1 < 1 || i2 > nx || i1 > i2) {
        cpl_msg_warning(__func__, "slice %2d: faulty polynomial detected at "
                        "y=%d (borders: %d ... %d)", islice + 1, ny/2, i1, i2);
        muse_trace_polys_delete(ptrace);
        continue; /* next slice */
      }
    } else { /* find slice edge without tracing info */
      char *keyword = cpl_sprintf("ESO DRS MUSE SLICE%d CENTER", (int)islice + 1);
      istart = cpl_propertylist_get_float(aImage->header, keyword);
      cpl_free(keyword);
      /* move outwards */
      int ix = istart;
      while (!(dq[(--ix-1) + ny/2*nx] & EURO3D_MISSDATA)) {
        i1 = ix;
      }
      ix = istart;
      while (!(dq[(++ix-1) + ny/2*nx] & EURO3D_MISSDATA)) {
        i2 = ix;
      }
    } /* else, no aTrace */
#if SKY_ROWBYROW_DEBUG
    cpl_msg_debug(__func__, "1 slice %d row %d edges: %d %d", islice+1, ny/2, i1, i2);
#endif
    /* loop through all image rows */
    int j;
    for (j = 0; j < ny; j++) {
      /* include pixels that have more than half of them inside the slice,   *
       * use the same method to compute the cutoff as muse_pixtable_create() */
      int ileft = istart, icenter = istart, iright = istart;
      if (ptrace) { /* slice edges with tracing info */
        ileft = ceil(cpl_polynomial_eval_1d(ptrace[MUSE_TRACE_LEFT], j+1, NULL)),
        icenter = cpl_polynomial_eval_1d(ptrace[MUSE_TRACE_CENTER], j+1, NULL),
        iright = floor(cpl_polynomial_eval_1d(ptrace[MUSE_TRACE_RIGHT], j+1, NULL));
      } else { /* find slice edges here without tracing info */
        /* move outwards */
        int ix = istart;
        while (!(dq[(--ix-1) + j*nx] & EURO3D_MISSDATA)) {
          ileft = ix;
        }
        ix = istart;
        while (!(dq[(++ix-1) + j*nx] & EURO3D_MISSDATA)) {
          iright = ix;
        }
        icenter = (iright + ileft) / 2.; /* re-"compute" center at this pos */
      } /* else, no aTrace */
#if SKY_ROWBYROW_DEBUG
      cpl_msg_debug(__func__, "2 slice %d row %d edges: %d %d", islice+1, j+1,
                    ileft, iright);
#endif

      /* fill data into vectors, copy the matrix */
      cpl_matrix *pos = cpl_matrix_new(1, iright - ileft + 1);
      cpl_vector *values = cpl_vector_new(iright - ileft + 1);
                 /* unused: *errors = cpl_vector_new(iright - ileft + 1); */
      unsigned int nval = 0;
      int i;
      for (i = ileft; i <= iright; i++) { /* x pos start at 1 */
        cpl_matrix_set(pos, 0, i - ileft, i - icenter); /* relative coordinates */
        if (dq[(i-1) + j*nx] != EURO3D_GOODPIXEL) {
          cpl_vector_set(values, i - ileft, NAN); /* mark as bad entry */
        } else {
          cpl_vector_set(values, i - ileft, data[(i-1) + j*nx]);
          nval++;
        }
        /* XXX sigmas are not supported by cpl_polynomial_fit() yet *
         * cpl_vector_set(errors, i - ileft, sqrt(data[i + j*nx])); */
      } /* for i (relevant columns) */
      if (nval < 1) { /* skip image rows with only NANs */
        cpl_vector_delete(values);
        /* unused: cpl_vector_delete(errors); */
        cpl_matrix_delete(pos);
        continue;
      }
      /* iterate the horizontal polynomial fit (originally 4th order) */
      unsigned int order = aOrder > nval + 1 ? nval - 1 : aOrder;
      double mse;
      cpl_polynomial *poly = muse_utils_iterate_fit_polynomial(pos, values,
                                                               NULL, NULL, order,
                                                               aRSigma, &mse,
                                                               NULL);
      int npix = cpl_vector_get_size(values); /* keep number of final pixels */
      cpl_vector_delete(values);
      /* unused: cpl_vector_delete(errors); */
      cpl_matrix_delete(pos);

      /* subtract the fit, add the mean-squared error onto the variance */
      for (i = ileft - 1; i < iright; i++) { /* x pos start at 0! */
        /* evaluate polynomial at the same relative position within slice */
        double sky = cpl_polynomial_eval_1d(poly, i+1 - icenter, NULL);
#if SKY_ROWBYROW_DEBUG > 1
        if (islice+1 == SKY_ROWBYROW_DEBUG_SLICE &&
            j+1 >= SKY_ROWBYROW_DEBUG_ROW1 && j+1 <= SKY_ROWBYROW_DEBUG_ROW2) {
          printf("subtracting slice %d row %d   %d %f  ", islice+1, j+1,
                 i+1 - icenter, data[i + j*nx]);
        }
#endif
        data[i + j*nx] -= sky;
        /* to correct the variance, add the mean-squared error divided *
         * by number of pixels minus the degrees of freedom            */
        stat[i + j*nx] += mse / (npix - order-1);
#if SKY_ROWBYROW_DEBUG > 1
        if (islice+1 == SKY_ROWBYROW_DEBUG_SLICE &&
            j+1 >= SKY_ROWBYROW_DEBUG_ROW1 && j+1 <= SKY_ROWBYROW_DEBUG_ROW2) {
          printf("%f\n", data[i + j*nx]);
          fflush(stdout);
        }
#endif
      } /* for i (relevant columns) */
      cpl_polynomial_delete(poly);
    } /* for j (all image rows) */

    muse_trace_polys_delete(ptrace); /* NULL check in there... */
  } /* for islice */

#if SKY_ROWBYROW_DEBUG
  cpl_msg_debug(__func__, "saving \"sky_removed_iterated_edge2.fits\"");
  muse_image_save(aImage, "sky_removed_iterated_edge2.fits");
#endif

  return CPL_ERROR_NONE;
} /* muse_sky_subtract_rowbyrow() */

/**@}*/
