/* -*- 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 <string.h>

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

/*---------------------------------------------------------------------------*
 *                             Functions code                                *
 *---------------------------------------------------------------------------*/

/*----------------------------------------------------------------------------*/
/**
  @brief    Add information for simplified wavelength solution to header.
  @param    aHeader   a FITS header to add the info to
  @param    aWC       the table with the real wavelength solution
  @param    aTT       the table with the tracing solution

  @error{set CPL_ERROR_NULL_INPUT and message and return,
         one of aHeader\, aWC\, or aTT is NULL}
 */
/*----------------------------------------------------------------------------*/
static void
muse_qi_mask_wavecal_polys(cpl_propertylist *aHeader, cpl_table *aWC,
                           cpl_table *aTT)
{
  if (!aHeader || !aWC || !aTT) {
    cpl_error_set_message(__func__, CPL_ERROR_NULL_INPUT, "Could not create per"
                          "-slice wavelength polynomials in the FITS header!");
    return;
  }

  /* number of evaluation steps for the wavelength calibration solution */
  const unsigned int kSteps = 50;

  /* create matrix and vector for the fitting process */
  cpl_matrix *pos = cpl_matrix_new(1, kMuseSlicesPerCCD * kSteps);
  cpl_vector *val = cpl_vector_new(kMuseSlicesPerCCD * kSteps),
             *ycenters = cpl_vector_new(kMuseSlicesPerCCD);
  cpl_size ipos = 0;
  unsigned short nslice;
  for (nslice = 1; nslice <= kMuseSlicesPerCCD; nslice++) {
    cpl_polynomial *pw = muse_wave_table_get_poly_for_slice(aWC, nslice),
                   **pt = muse_trace_table_get_polys_for_slice(aTT, nslice);
    /* evaluate central trace at vertical center of the CCD */
    double x = cpl_polynomial_eval_1d(pt[MUSE_TRACE_CENTER], kMuseOutputYTop/2,
                                      NULL);
    muse_trace_polys_delete(pt);

    cpl_polynomial *pconst = cpl_polynomial_new(1);
    cpl_size pows = 0;
    cpl_polynomial_set_coeff(pconst, &pows, x);
    /* reduce the polynomial to 1D at the center of the slice */
    cpl_polynomial *pwcen = cpl_polynomial_extract(pw, 0, pconst);
    double ycen = cpl_polynomial_eval_1d(pwcen, kMuseOutputYTop/2, NULL);
    cpl_vector_set(ycenters, nslice - 1, ycen);
    cpl_polynomial_delete(pconst);
    cpl_polynomial_delete(pw);

    /* shift the polynomial to a common center */
    pows = 0;
    double c0 = cpl_polynomial_get_coeff(pwcen, &pows);
    cpl_polynomial_set_coeff(pwcen, &pows, c0 - ycen);

    unsigned int iy;
    for (iy = 1;
         iy <= (unsigned)kMuseOutputYTop;
         iy += ((unsigned)kMuseOutputYTop) / (kSteps - 1)) {
      cpl_matrix_set(pos, 0, ipos, iy);
      cpl_vector_set(val, ipos++, cpl_polynomial_eval_1d(pwcen, iy, NULL));
    }
  } /* for nslice */

  /* now try to fit the common 2nd-order polynomial */
  cpl_polynomial *fit = cpl_polynomial_new(1);
  cpl_size maxdeg = 2;
  cpl_polynomial_fit(fit, pos, NULL, val, NULL, CPL_TRUE, NULL, &maxdeg);
#if 0 /* DEBUG */
  cpl_vector *res = cpl_vector_duplicate(val);
  double chisq;
  cpl_vector_fill_polynomial_fit_residual(res, val, NULL, fit, pos, &chisq);
  const double mse = cpl_vector_product(res, res)
                   / cpl_vector_get_size(res);
  cpl_vector_delete(res);
  cpl_msg_debug(__func__, "MSE=%g ChiSq=%g", mse, chisq);
#endif
  cpl_vector_delete(val);
  cpl_matrix_delete(pos);

  /* now finally add the relevant keywords to the header for this IFU */
  char keyword[KEYWORD_LENGTH];
  cpl_size j = 0; /* accessor to polynomial coefficients */
  for (nslice = 1; nslice <= kMuseSlicesPerCCD; nslice++) {
    snprintf(keyword, KEYWORD_LENGTH, "ESO DET WLEN SLICE%hu POLY0", nslice);
    /* add the vertical shift back again */
    double zeropoint = cpl_polynomial_get_coeff(fit, &j)
                     + cpl_vector_get(ycenters,  nslice - 1);
    cpl_propertylist_append_float(aHeader, keyword, zeropoint);
    cpl_propertylist_set_comment(aHeader, keyword,
                                 "[Angstrom] Slice-specific zero-point");
  } /* for nslice */
  cpl_vector_delete(ycenters);

  /* also add the common coefficients for 1st and 2nd order */
  for (j = 1; j <= 2; j++) {
    snprintf(keyword, KEYWORD_LENGTH, "ESO DET WLEN POLY%d", (int)j);
    cpl_propertylist_append_float(aHeader, keyword,
                                  cpl_polynomial_get_coeff(fit, &j));
    char comment[100];
    snprintf(comment, 99, "[Angstrom] Common %s-order polynomial coefficient",
             j == 1 ? "1st" : "2nd");
    cpl_propertylist_set_comment(aHeader, keyword, comment);
  } /* for j (polynomial orders) */
  cpl_polynomial_delete(fit);
} /* muse_qi_mask_wavecal_polys() */

/*----------------------------------------------------------------------------*/
/**
  @brief    Interpret the command line options and execute the data processing
  @param    aProcessing  the processing structure
  @param    aParams      the parameters list
  @return   0 if everything is ok, -1 something went wrong

  Compared to other recipes, muse_qi_mask is special-cased in
  xmldoc/recipedoc2c.xsl so that aParams->nifu stays zero if zero was passed to
  the recipe.
 */
/*----------------------------------------------------------------------------*/
int
muse_qi_mask_compute(muse_processing *aProcessing,
                     muse_qi_mask_params_t *aParams)
{
  /* check for wrange and/or wmin/wmax parameters */
  double wmin = 4650.,
         wmax = 9300.;

  cpl_msg_info(__func__, "Creating mask for full MUSE wavelength range "
               "(%.1f...%.1f)", wmin, wmax);

  /* if we find a BIAS input, add it to usedframes to get it *
   * to be used for creation of the DFS-compliant header     */
  cpl_frame *frame = cpl_frameset_find(aProcessing->inframes, MUSE_TAG_BIAS);
  if (frame) {
    muse_processing_append_used(aProcessing, frame, CPL_FRAME_GROUP_RAW, 1);
  }

  /* create at least minimal header and save it */
  cpl_propertylist *header = cpl_propertylist_new();
  cpl_propertylist_append_string(header, "OBJECT", "Mask full MUSE range");
  cpl_propertylist_append_float(header, "ESO DET FRS WMIN", wmin);
  cpl_propertylist_set_comment(header, "ESO DET FRS WMIN", "[Angstrom] Minimum wavelength");
  cpl_propertylist_append_float(header, "ESO DET FRS WMAX", wmax);
  cpl_propertylist_set_comment(header, "ESO DET FRS WMAX", "[Angstrom] Maximum wavelength");
  muse_processing_save_header(aProcessing, aParams->nifu ? aParams->nifu : -1,
                              header, MUSE_TAG_MASK_IMAGE);
  cpl_propertylist_delete(header);

  /* get output filename from output frameset */
  frame = cpl_frameset_find(aProcessing->outframes, MUSE_TAG_MASK_IMAGE);
  const char *fn = cpl_frame_get_filename(frame);

  /* loop over all 24 IFUs */
  int n1 = aParams->nifu ? aParams->nifu : 1,
      n2 = aParams->nifu ? aParams->nifu : kMuseNumIFUs,
      nifu;
  for (nifu = n1; nifu <= n2; nifu++) {
    /* load raw (bias) image and trim it or create minimal default image */
    cpl_errorstate prestate = cpl_errorstate_get();
    muse_imagelist *images = muse_basicproc_load(aProcessing, nifu, NULL);
    muse_image *image = NULL;
    if (images) {
      cpl_msg_debug(__func__, "succeeded to load %u raw images, using first one",
                    muse_imagelist_get_size(images));
      image = muse_imagelist_get(images, 0);
    } else {
      cpl_msg_warning(__func__, "failed to load raw image, assuming 1x1 binning");
      cpl_errorstate_dump(prestate, CPL_FALSE, muse_cplerrorstate_dump_some);
      image = muse_image_new();
      /* use typical 1x1 binned trimmed MUSE image size */
      image->data = cpl_image_new(kMuseOutputXRight, kMuseOutputYTop,
                                  CPL_TYPE_FLOAT);
      image->header = cpl_propertylist_new();
      cpl_propertylist_append_int(image->header, "ESO DET BINX", 1);
      cpl_propertylist_append_int(image->header, "ESO DET BINY", 1);
      images = muse_imagelist_new();
      muse_imagelist_set(images, image, 0); /* to easily delete it below */
    }

    /* load wavecaltable and tracetable */
    cpl_table *wavecaltable = muse_processing_load_ctable(aProcessing,
                                                          MUSE_TAG_WAVECAL_TABLE,
                                                          nifu),
              *tracetable = muse_processing_load_ctable(aProcessing,
                                                        MUSE_TAG_TRACE_TABLE, nifu);

    /* create the wavelength map, first always unbinned */
    int binx = muse_pfits_get_binx(image->header),
        biny = muse_pfits_get_biny(image->header);
    muse_image *imunbinned = image;
    if (binx != 1 || biny != 1) { /* need unbinned image for wavemep to work */
      imunbinned = muse_image_new();
      imunbinned->data = cpl_image_new(kMuseOutputXRight, kMuseOutputYTop,
                                       CPL_TYPE_FLOAT);
    }
    cpl_image *wavemap = muse_wave_map(imunbinned, wavecaltable, tracetable);

    /* now take into account the binning */
    if (binx != 1 || biny != 1) {
      muse_image_delete(imunbinned);
      cpl_msg_info(__func__, "Rebinning wavelength map %dx%d", binx, biny);
      cpl_image *binned = cpl_image_rebin(wavemap, 1, 1, binx, biny);
      cpl_image_delete(wavemap);
      /* cpl_image_rebin() creates summed image, we need averaged */
      wavemap = cpl_image_divide_scalar_create(binned, binx * biny);
      cpl_image_delete(binned);
    }

    /* apply the threshold in wavelengths */
    cpl_mask *mask = cpl_mask_threshold_image_create(wavemap, wmin, wmax);

    /* create the output FITS header, for properties in this extension */
    header = cpl_propertylist_new();
    if (binx == 1 && biny == 1) {
      muse_qi_mask_wavecal_polys(header, wavecaltable, tracetable);
    }

    /* Analyze mask to determine slice positions and store them in the  *
     * header. Use the trimmed mask, so that cpl_apertures_get_size()   *
     * finds 48 connected regions instead of 96 with the overscan gap.  *
     * Also make up for the pre- and overscans by adding n*32/bin.      */
    int nx = cpl_mask_get_size_x(mask),
        ny = cpl_mask_get_size_y(mask);
    cpl_size ndet;
    cpl_image *labeled = cpl_image_labelise_mask_create(mask, &ndet);
    /* it is simplest to use apertures to get the coordinates that we want */
    cpl_apertures *apertures = cpl_apertures_new_from_image(labeled, labeled);
    int napertures = cpl_apertures_get_size(apertures);
    /* the apertures are not sorted by x position and we cannot    *
     * sort apertures externally to CPL, so use an x-sorted object *
     * to access slice numbers directly                            */
    cpl_matrix *mlabels = cpl_matrix_new(2, napertures);
    int n; /* aperture number = label number */
    for (n = 1; n <= napertures; n++) {
      cpl_matrix_set(mlabels, 0, n - 1,
                     cpl_apertures_get_left(apertures, n)); /* x-position in row. 0 */
      cpl_matrix_set(mlabels, 1, n - 1, n); /* aperture number in row. 1 */
    } /* for n (all apertures) */
    cpl_matrix_sort_columns(mlabels, 1); /* sort by x-position value */
#if 0
    cpl_matrix_dump(mlabels, stdout);
#endif
    int m, /* index entry in matrix */
        nslice = 1; /* real slice number */
    for (m = 1; m <= ndet; m++) {
      n = cpl_matrix_get(mlabels, 1, m - 1); /* get aperture number from sorted matrix */
      /* aperture extremes, with prescan added */
      int x1 = cpl_apertures_get_left(apertures, n),
          x2 = cpl_apertures_get_right(apertures, n),
          y1 = cpl_apertures_get_bottom(apertures, n),
          y2 = cpl_apertures_get_top(apertures, n);
      if (x1 == x2 || y1 == y2) {
        /* need to exclude single-pixel apertures, created by rebinning; *
         * also replace them with zeros in the mask */
        int i;
        for (i = x1; i <= x2; i++) {
          int j;
          for (j = y1; j <= y2; j++) {
            cpl_mask_set(mask, i, j, CPL_BINARY_0);
          } /* for j (all y pixels of detection) */
        } /* for i (all x pixels of detection) */
        cpl_msg_debug(__func__, "Excluding aperture [%d:%d,%d:%d]", x1, x2, y1, y2);
        continue;
      }
      /* add prescan to aperture extremes */
      x1 += 32 / binx;
      x2 += 32 / binx;
      y1 += 32 / biny;
      y2 += 32 / biny;
      /* add overscan gaps */
      if (x1 > nx/2+32/binx) x1 += 64 / binx;
      if (x2 > nx/2+32/binx) x2 += 64 / binx;
      if (y1 > ny/2+32/biny) y1 += 64 / biny;
      if (y2 > ny/2+32/biny) y2 += 64 / biny;
#if 0
      int xc = (x1 + x2) / 2,
          yc = (y1 + y2) / 2,
          w = x2 - x1 + 1,
          h = y2 - y1 + 1;
      cpl_msg_debug(__func__, "Aperture %2d Slice %2d: [%d:%d,%d:%d] center %d,%d size %dx%d",
                    n, m, x1, x2, y1, y2, xc, yc, w, h);
#endif
      char keyword[KEYWORD_LENGTH];
      snprintf(keyword, KEYWORD_LENGTH, "ESO DET SLICE%d XSTART", nslice);
      cpl_propertylist_append_int(header, keyword, x1);
      cpl_propertylist_set_comment(header, keyword, "[pix] Start position of the slice along X");
      snprintf(keyword, KEYWORD_LENGTH, "ESO DET SLICE%d YSTART", nslice);
      cpl_propertylist_append_int(header, keyword, y1);
      cpl_propertylist_set_comment(header, keyword, "[pix] Start position of the slice along Y");
      snprintf(keyword, KEYWORD_LENGTH, "ESO DET SLICE%d XEND", nslice);
      cpl_propertylist_append_int(header, keyword, x2);
      cpl_propertylist_set_comment(header, keyword, "[pix] End position of the slice along X");
      snprintf(keyword, KEYWORD_LENGTH, "ESO DET SLICE%d YEND", nslice);
      cpl_propertylist_append_int(header, keyword, y2);
      cpl_propertylist_set_comment(header, keyword, "[pix] End position of the slice along Y");
      nslice++;
    } /* for m (all matrix entries) */
    cpl_image_delete(labeled);
    cpl_apertures_delete(apertures);
    cpl_matrix_delete(mlabels);
    if (nslice - 1 != kMuseSlicesPerCCD) {
      cpl_msg_warning(__func__, "Found %d slices (%d apertures) instead of %d",
                      nslice - 1, napertures, kMuseSlicesPerCCD);
    }

    /* re-add pre- and overscans, assuming 32/bin pixels on all sides */
    cpl_mask *untrimmed = cpl_mask_new(nx + 4*32/binx, ny + 4*32/biny);
    cpl_mask *m1 = cpl_mask_extract(mask, 1, 1, nx/2-1, ny/2-1),
             *m2 = cpl_mask_extract(mask, 1, ny/2, nx/2-1, ny),
             *m3 = cpl_mask_extract(mask, nx/2, 1, nx, ny/2-1),
             *m4 = cpl_mask_extract(mask, nx/2, ny/2, nx, ny);
    cpl_mask_copy(untrimmed, m1, 32/binx+1, 32/biny+1);
    cpl_mask_copy(untrimmed, m2, 32/binx+1, ny/2+3*32/biny);
    cpl_mask_copy(untrimmed, m3, nx/2+3*32/binx, 32/biny+1);
    cpl_mask_copy(untrimmed, m4, nx/2+3*32/binx, ny/2+3*32/biny);
    cpl_mask_delete(m1);
    cpl_mask_delete(m2);
    cpl_mask_delete(m3);
    cpl_mask_delete(m4);
    cpl_mask_delete(mask);
    mask = untrimmed;

    /* save the image to extension CHAN%02d, as 8bit mask image */
    if (image->header) {
      /* copy relevant detector properties from input image */
      cpl_propertylist_copy_property_regexp(header, image->header,
                                            "ESO DET (CHIP |OUT)|EXTNAME", 0);
    } else {
      /* create minimal header with only EXTNAME */
      char *extname = cpl_sprintf("CHAN%02d", nifu);
      cpl_propertylist_append_string(header, "EXTNAME", extname);
      cpl_free(extname);
    }
    cpl_mask_save(mask, fn, header, CPL_IO_EXTEND);
    cpl_propertylist_delete(header);
    cpl_image_delete(wavemap);
    cpl_mask_delete(mask);
    cpl_table_delete(tracetable);
    cpl_table_delete(wavecaltable);
    muse_imagelist_delete(images);
  } /* for nifu */

  return 0;
} /* muse_qi_mask_compute() */
