/* -*- 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) 2014-2016 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 <muse.h>
#include <string.h>

/*----------------------------------------------------------------------------*/
/**
 * @defgroup tools_musecubecombine          Tool muse_cube_combine
 *
 * <b>muse_cube_combine</b>: Combine fully reduced MUSE cubes.
 *
 * Combine cubes that were resampled onto the same output grid but using data of
 * different wavelength ranges into a single output cube.
 *
 * <b>Command line arguments:</b>
 *   - <tt><b>CUBE_OUT</b></tt>\n
 *     the filename of output cube
 *   - <tt><b>CUBE_IN_1</b></tt>\n
 *     the filename of the 1st input cube to use for combination
 *   - <tt><b>CUBE_IN_2</b></tt>\n
 *     the filename of the 2nd input cube to use for combination
 *   - <tt>CUBE_IN_3 ...</tt> (optional)\n
 *     the filename(s) of the 3rd and following inputs cubes to use for
 *     combination; the number of total cubes is not limited, but at least
 *     two input cubes need to be given for this tool to make sense
 *
 * <b>Return values:</b>
 *   - <tt> 0</tt>\n   Success
 *   - <tt> 1</tt>\n   less than three filenames given
 *   - <tt> 9</tt>\n   unknown option given
 *   - <tt>10</tt>\n   no output cube name found
 *   - <tt>11</tt>\n   cube with given output filename already exists
 *   - <tt>12</tt>\n   less than 2 valid cubes found in input
 *   - <tt>20</tt>\n   error when saving the output cube
 *   - <tt>50</tt>\n   unknown error
 */
/*----------------------------------------------------------------------------*/

/**@{*/

#define PRINT_USAGE(rc)                                                        \
  fprintf(stderr, "Usage: %s CUBE_OUT CUBE_IN_1 CUBE_IN_2 [ CUBE_IN_3 ... ]\n",\
          argv[0]);                                                            \
  cpl_end(); return (rc);

int main(int argc, char **argv)
{
  const char *idstring = "muse_cube_combine";
  cpl_init(CPL_INIT_DEFAULT);
  cpl_msg_set_time_on();
  muse_processing_recipeinfo(NULL);
  cpl_msg_set_level(CPL_MSG_DEBUG);
  cpl_msg_set_component_on();
  cpl_errorstate state = cpl_errorstate_get();

  if (argc < 4) {
    /* three filenames are needed at least */
    PRINT_USAGE(1);
  }

  char *oname = NULL; /* output cube */
  int i, noverlap = 1;

  /* argument processing */
  for (i = 1; i < argc; i++) {
    if (strncmp(argv[i], "-xxx", 5) == 0) { // XXX
      /* skip to next arg to get start value */
      i++;
//      if (i < argc) {
//        x1 = atof(argv[i]);
//      } else {
//        PRINT_USAGE(2);
//      }
    } else if (strncmp(argv[i], "-", 1) == 0) { /* unallowed options */
      PRINT_USAGE(9);
    } else {
      if (oname) {
        break; /* we have the required name, skip the rest */
      }
      oname = argv[i];
    }
  } /* for i (all arguments) */
  if (!oname) {
    PRINT_USAGE(10);
  }
  FILE *fp = fopen(oname, "r");
  if (fp) {
    cpl_msg_error(idstring, "Output cube \"%s\" is already present!", oname);
    cpl_msg_error(idstring, "Please specify another output name or rename the "
                  "existing file with this name.");
    fclose(fp);
    PRINT_USAGE(11);
  }

  int ncubes = argc - i;
  cpl_msg_info(idstring, "Will write combination of %d cubes to \"%s\".",
               ncubes, oname);

  /* create table of the cubes involved here */
  cpl_table *table = cpl_table_new(ncubes);
  cpl_table_new_column(table, "FILENAME", CPL_TYPE_STRING);
  cpl_table_new_column(table, "LMIN", CPL_TYPE_DOUBLE);
  cpl_table_new_column(table, "LMAX", CPL_TYPE_DOUBLE);
  cpl_table_new_column(table, "P1", CPL_TYPE_INT);
  cpl_table_new_column(table, "P2", CPL_TYPE_INT);

  /* check ranges and grids of the cubes involved, save in the table */
  int irecout = -1, /* output receipe parameter set */
      iparout1 = -1, /* and output parameter numbers for low */
      iparout2 = -1; /* and high wavelength limits */
  double lminout = 99999.,  /* variables to store final output limit for low */
         lmaxout = -99999.; /* and high wavelength limits */
  char *ctype1ref = NULL, *ctype2ref = NULL,
       *cunit1ref = NULL, *cunit2ref = NULL;
  double crpix1ref = NAN, crpix2ref = NAN, crval1ref = NAN, crval2ref = NAN,
         cd11ref = NAN, cd12ref = NAN, cd21ref = NAN, cd22ref = NAN;
  char *ctype3ref = NULL;
  cpl_boolean loglambda = CPL_FALSE;
  double crpix3ref = NAN,
         crval3ref = NAN,
         cd33ref = NAN;
  int ic;
  for (ic = 0; ic < ncubes; ic++) {
    int argidx = ic + 2;
    char *iname = argv[argidx];
    /* search for the data extension */
    int iext = cpl_fits_find_extension(iname, "DATA");
    if (iext < 0) {
      cpl_msg_warning(idstring, "Could not open cube \"%s\"", iname);
      continue;
    }

    /* load primary header and get lambda ranges */
    cpl_propertylist *header = cpl_propertylist_load(iname, 0);
    double lmin = -FLT_MAX,
           lmax = FLT_MAX;
    /* search for last set of recipe parameters stored in the cube,  */
    cpl_errorstate es = cpl_errorstate_get();
    int irec = 0;
    do {
      irec++;
      char *hname = cpl_sprintf("ESO PRO REC%d ID", irec);
      /* recipe name should exist at least, so search that */
      const char *name = cpl_propertylist_get_string(header, hname);
      cpl_free(hname);
      if (!name) {
        break;
      }
      cpl_msg_debug(idstring, "REC = %d, ID = %s", irec, name);
    } while (cpl_errorstate_is_equal(es));
    cpl_errorstate_set(es);
    cpl_msg_info(idstring, "Using recipe parameters ESO PRO REC%d to determine "
                 "wavelength ranges", --irec);
    if (irecout < 0) {
       irecout = irec; /* remember output receipe parameter set */
    }

    /* loop through the recipe parameters stored in the cube */
    int ipar = 0;
    do {
      ipar++;
      char *hname = cpl_sprintf("ESO PRO REC%d PARAM%d NAME", irec, ipar);
      const char *name = cpl_propertylist_get_string(header, hname);
      cpl_free(hname);
      cpl_boolean islmin = name && !strncmp(name, "lambdamin", 10),
                  islmax = name && !strncmp(name, "lambdamax", 10);
      if (islmin || islmax) {
        char *hval = cpl_sprintf("ESO PRO REC%d PARAM%d VALUE", irec, ipar);
        if (islmin) {
          if (iparout1 < 0) {
            iparout1 = ipar; /* remember parameter number */
          }
          lmin = atof(cpl_propertylist_get_string(header, hval));
          cpl_msg_debug(idstring, "lmin = %f", lmin);
          if (lmin < lminout) {
            lminout = lmin; /* remember minimum output limit */
          }
        } else {
          if (iparout2 < 0) {
            iparout2 = ipar; /* remember parameter number */
          }
          lmax = atof(cpl_propertylist_get_string(header, hval));
          cpl_msg_debug(idstring, "lmax = %f", lmax);
          if (lmax > lmaxout) {
            lmaxout = lmax; /* remember maximum output limit */
          }
        }
        cpl_free(hval);
      } /* if */
    } while (cpl_errorstate_is_equal(es));
    cpl_errorstate_set(es);
    cpl_propertylist_delete(header);

    /* load DATA header and get lambda solution */
    header = cpl_propertylist_load(iname, iext);
    double crpix3 = muse_pfits_get_crpix(header, 3),
           crval3 = muse_pfits_get_crval(header, 3),
           cd33 = muse_pfits_get_cd(header, 3, 3);
    int naxis3 = cpl_propertylist_get_int(header, "NAXIS3");

    const char *cunit3 = muse_pfits_get_cunit(header, 3),
               *lunit = cpl_table_get_column_unit(table, "LMIN");
    if (!lunit) {
      cpl_table_set_column_unit(table, "LMIN", cunit3);
      cpl_table_set_column_unit(table, "LMAX", cunit3);
    } else if (strncmp(lunit, cunit3, strlen(lunit) + 1)) {
      cpl_msg_warning(idstring, "Cube %d does not contain wavelengths in "
                      "%s (%s), skipping it", ic + 1, lunit, cunit3);
      cpl_propertylist_delete(header);
      continue;
    }
    const char *ctype3 = muse_pfits_get_ctype(header, 3);
    if (!ctype3ref) {
      ctype3ref = cpl_strdup(ctype3);
      crpix3ref = crpix3;
      crval3ref = crval3;
      cd33ref = cd33;
      loglambda = ctype3ref && (!strncmp(ctype3ref, "AWAV-LOG", 9) ||
                                !strncmp(ctype3ref, "WAVE-LOG", 9));
    } else if (strncmp(ctype3ref, ctype3, strlen(ctype3ref)) ||
               fabs(crpix3 - crpix3ref) > DBL_EPSILON ||
               fabs(crval3 - crval3ref) > DBL_EPSILON ||
               fabs(cd33 - cd33ref) > DBL_EPSILON) {
      cpl_msg_warning(idstring, "Cube %d does not match in spectral WCS, "
                      "skipping it", ic + 1);
      cpl_propertylist_delete(header);
      continue;
    }
    double lpx1 = (1. - crpix3) * cd33 + crval3,
           lpx2 = (naxis3 - crpix3) * cd33 + crval3;
    if (loglambda) {
      lpx1 = crval3 * exp((1 - crpix3) * cd33 / crval3);
      lpx2 = crval3 * exp((naxis3 - crpix3) * cd33 / crval3);
    }
    cpl_msg_debug(idstring, "Cube %d, axis 3 WCS: %f %f %f (%s, %d pixels, "
                  "%f..%f %s)", ic + 1, crpix3, crval3, cd33, ctype3, naxis3,
                  lpx1, lpx2, cunit3);
    /* check min/max wavelengths against boundary conditions */
    if (lmin < lpx1) {
      lmin = lpx1;
    }
    if (lmax > lpx2) {
      lmax = lpx2;
    }

    /* check spatial WCS against first cube */
    const char *ctype1 = muse_pfits_get_ctype(header, 1),
               *ctype2 = muse_pfits_get_ctype(header, 2),
               *cunit1 = muse_pfits_get_cunit(header, 1),
               *cunit2 = muse_pfits_get_cunit(header, 2);
    double crpix1 = muse_pfits_get_crpix(header, 1),
           crpix2 = muse_pfits_get_crval(header, 2),
           crval1 = muse_pfits_get_crval(header, 1),
           crval2 = muse_pfits_get_crval(header, 2),
           cd11 = muse_pfits_get_cd(header, 1, 1),
           cd12 = muse_pfits_get_cd(header, 1, 2),
           cd21 = muse_pfits_get_cd(header, 2, 1),
           cd22 = muse_pfits_get_cd(header, 2, 2);
    if (!ctype1ref) {
      ctype1ref = cpl_strdup(ctype1);
      ctype2ref = cpl_strdup(ctype2);
      cunit1ref = cpl_strdup(cunit1);
      cunit2ref = cpl_strdup(cunit2);
      crpix1ref = crpix1;
      crval1ref = crval1;
      crpix2ref = crpix2;
      crval2ref = crval2;
      cd11ref = cd11;
      cd12ref = cd12;
      cd21ref = cd21;
      cd22ref = cd22;
    } else if (strncmp(ctype1ref, ctype1, strlen(ctype1ref)) ||
               strncmp(ctype2ref, ctype2, strlen(ctype2ref)) ||
               strncmp(cunit1ref, cunit1, strlen(cunit1ref)) ||
               strncmp(cunit2ref, cunit2, strlen(cunit2ref)) ||
               fabs(crpix1 - crpix1ref) > DBL_EPSILON ||
               fabs(crval1 - crval1ref) > DBL_EPSILON ||
               fabs(crpix2 - crpix2ref) > DBL_EPSILON ||
               fabs(crval2 - crval2ref) > DBL_EPSILON ||
               fabs(cd11 - cd11ref) > DBL_EPSILON ||
               fabs(cd12 - cd12ref) > DBL_EPSILON ||
               fabs(cd21 - cd21ref) > DBL_EPSILON ||
               fabs(cd22 - cd22ref) > DBL_EPSILON) {
      cpl_msg_warning(idstring, "Cube %d does not match in spatial WCS, "
                      "skipping it", ic + 1);
      cpl_propertylist_delete(header);
      continue;
    }
    cpl_propertylist_delete(header);

    /* compute good valid range of pixels in dispersion direction */
    int l1 = lround((lmin - crval3) / cd33 + crpix3),
        l2 = lround((lmax - crval3) / cd33 + crpix3);
    if (loglambda) {
      l1 = lround(crval3 / cd33 * log(lmin / crval3)) + crpix3;
      l2 = lround(crval3 / cd33 * log(lmax / crval3)) + crpix3;
    }
    cpl_msg_debug(idstring, "Cube %d: %d..%d (%f..%f Angstrom)", ic + 1,
                  l1, l2, lmin, lmax);
    if (l1 <= 1) {
      l1 = 1;
    } else if (l1 >= naxis3) {
      cpl_msg_warning(idstring, "Something is wrong with cube %d: l1 = %d!",
                      ic + 1, l1);
      continue;
    } else {
      /* throw away plane(s) from the starting edge to mitigate resampling edge effects */
      l1 += noverlap;
    }
    if (l2 >= naxis3) {
      l2 = naxis3;
    } else if (l2 <= 1) {
      cpl_msg_warning(idstring, "Something is wrong with cube %d: l2 = %d!",
                      ic + 1, l2);
      continue;
    } else {
      l2 -= noverlap; /* throw away plane(s) from the ending edge */
    }
    cpl_msg_debug(idstring, "Cube %d: ---> %d..%d", ic + 1, l1, l2);

    /* all was OK, so we can save info of this cube to the table */
    cpl_table_set_string(table, "FILENAME", ic, iname);
    cpl_table_set(table, "LMIN", ic, lmin);
    cpl_table_set(table, "LMAX", ic, lmax);
    cpl_table_set(table, "P1", ic, l1);
    cpl_table_set(table, "P2", ic, l2);
  } /* for ic (all input cubes) */
  cpl_free(ctype1ref);
  cpl_free(ctype2ref);
  cpl_free(cunit1ref);
  cpl_free(cunit2ref);
  cpl_free(ctype3ref);

  /* erase bad entries (cubes with missing properties) */
  cpl_table_select_all(table);
  cpl_table_and_selected_invalid(table, "FILENAME");
  cpl_table_erase_selected(table);
  ncubes = cpl_table_get_nrow(table);
  if (ncubes < 2) {
    cpl_msg_error(idstring, "Only %d cube%s with valid info %s found!", ncubes,
                  ncubes ? "" : "s", ncubes ? "was" : "were");
    cpl_end();
    return 12;
  }

  /* now sort the table by increasing first plane */
  cpl_propertylist *order = cpl_propertylist_new();
  cpl_propertylist_append_bool(order, "P1", CPL_FALSE);
  cpl_table_sort(table, order);
  cpl_propertylist_delete(order);
#if 1
  printf("Stored and sorted %d cubes:\n", ncubes);
  cpl_table_dump(table, 0, ncubes, stdout);
  fflush(stdout);
#endif

  /* clean up overlaps and gaps:                                     *
   * - if there are overlaps just move the 2nd cube to start at a    *
   *   higher plane                                                  *
   * - if there is a gap just move the 1st cube to be more extended, *
   *   even though this will incur resampling edge effects (warn!)   */
  int nover = 0, ngaps = 0;
  for (ic = 1; ic < cpl_table_get_nrow(table); ic++) { /* start at 2nd row */
    int p2a = cpl_table_get_int(table, "P2", ic - 1, NULL),
        p1b = cpl_table_get_int(table, "P1", ic, NULL);
    double lmaxa = cpl_table_get_double(table, "LMAX", ic - 1, NULL);
    int pdiff = p1b - p2a;
    if (pdiff == 1) { /* no gap, no overlap */
      continue;
    }
    if (pdiff < 1) { /* overlap */
      p2a += pdiff - 1;
      lmaxa += cd33ref * (pdiff - 1);
      nover++;
    } /* if overlap */
    if (pdiff > 1) { /* gap */
      p2a += pdiff - 1;
      lmaxa += cd33ref * (pdiff - 1);
      ngaps++;
    } /* if gap */
    cpl_table_set_int(table, "P2", ic - 1, p2a);
    cpl_table_set_double(table, "LMAX", ic - 1, lmaxa);
  } /* for ic (all other valid cubes) */
#if 1
  printf("Cleaned up %d overlaps and %d gaps.\n", nover, ngaps);
  cpl_table_dump(table, 0, ncubes, stdout);
  fflush(stdout);
#endif

  /* now create the initial cube from the first exposure in the table */
  muse_datacube *cube = cpl_calloc(1, sizeof(muse_datacube));
  const char *fn = cpl_table_get_string(table, "FILENAME", 0);
  cpl_msg_info(idstring, "Loading cube from \"%s\"", fn);
  cube->header = cpl_propertylist_load(fn, 0);
  /* find DATA extension */
  int iext = cpl_fits_find_extension(fn, "DATA");
  /* load the WCS from the first cube extension *
   * and merge it into the cube header          */
  cpl_propertylist *wcs = cpl_propertylist_load(fn, iext);
  cpl_propertylist_erase_regexp(wcs, MUSE_WCS_KEYS, 1); /* remove others */
  cpl_propertylist_append(cube->header, wcs);
  cpl_propertylist_delete(wcs);
  /* load the cubes from the DATA and STAT extensions */
  cube->data = cpl_imagelist_load(fn, CPL_TYPE_FLOAT, iext);
  iext = cpl_fits_find_extension(fn, "STAT");
  cube->stat = cpl_imagelist_load(fn, CPL_TYPE_FLOAT, iext);

  /* now create flags vector (starting from array) to erase unused planes */
  int nplanes = cpl_imagelist_get_size(cube->data),
      l1 = cpl_table_get_int(table, "P1", 0, NULL),
      l2 = cpl_table_get_int(table, "P2", 0, NULL);
  cpl_array *aflags = cpl_array_new(nplanes, CPL_TYPE_DOUBLE);
  cpl_array_fill_window_double(aflags, 0, nplanes, -1.); /* all bad */
  cpl_array_fill_window_double(aflags, l1 - 1, l2 - l1 + 1, 1.); /* the good ones */
#if 0
  cpl_array_dump(aflags, 0, 5, stdout);
  cpl_array_dump(aflags, l1 - 4, 10, stdout);
  cpl_array_dump(aflags, l2 - 4, 10, stdout);
  cpl_array_dump(aflags, nplanes - 5, 6, stdout);
  fflush(stdout);
#endif
  cpl_vector *vflags = cpl_vector_wrap(nplanes, cpl_array_unwrap(aflags));
  cpl_imagelist_erase(cube->data, vflags);
  cpl_imagelist_erase(cube->stat, vflags);
  cpl_vector_delete(vflags);
  cpl_msg_debug(idstring, "%d of %d planes left in first cube",
                (int)cpl_imagelist_get_size(cube->data), nplanes);
  nplanes = cpl_imagelist_get_size(cube->data);
  /* update the spectral part of the WCS with the new setup, i.e. *
   * adjust CRPIX3 by the number of deleted planes at the start   */
  cpl_propertylist_update_double(cube->header, "CRPIX3", crpix3ref - (l1 - 1));

  /* one-by-one fully load the other cubes, and copy *
   * the valid data from them into the output cube   */
  for (ic = 1; ic < cpl_table_get_nrow(table); ic++) {
    fn = cpl_table_get_string(table, "FILENAME", ic);
    cpl_msg_info(idstring, "Loading cube from \"%s\"", fn);
    iext = cpl_fits_find_extension(fn, "DATA");
    cpl_imagelist *cubedata = cpl_imagelist_load(fn, CPL_TYPE_FLOAT, iext);
    iext = cpl_fits_find_extension(fn, "STAT");
    cpl_imagelist *cubestat = cpl_imagelist_load(fn, CPL_TYPE_FLOAT, iext);

    /* simply append (copy) the valid image planes to the output cube imagelists */
    l1 = cpl_table_get_int(table, "P1", ic, NULL);
    l2 = cpl_table_get_int(table, "P2", ic, NULL);
    int l;
    for (l = l1 - 1; l < l2; l++, nplanes = cpl_imagelist_get_size(cube->data)) {
      cpl_error_code rc1, rc2;
      rc1 = cpl_imagelist_set(cube->data,
                              cpl_image_duplicate(cpl_imagelist_get(cubedata, l)),
                              nplanes);
      rc2 = cpl_imagelist_set(cube->stat,
                              cpl_image_duplicate(cpl_imagelist_get(cubestat, l)),
                              nplanes);
      if (rc1 != rc2 || rc1 != CPL_ERROR_NONE || rc2 != CPL_ERROR_NONE) {
        cpl_msg_warning(idstring, "Could not append plane %d of cube %d to output cube",
                        l, ic + 1);
      } /* if bad result */
    } /* for l (wavelength planes to copy) */
    cpl_imagelist_delete(cubedata);
    cpl_imagelist_delete(cubestat);
    cpl_msg_debug(idstring, "output cube now has %d/%d planes after copying %d "
                  "planes from cube %d", nplanes,
                  (int)cpl_imagelist_get_size(cube->data), l2 - l1 + 1, ic + 1);
  } /* for ic (all other valid cubes) */
  cpl_table_delete(table);

  /* reset the recipe parameters in the output header with the now full range */
  cpl_msg_info(idstring, "Resetting recipe parameters ESO PRO REC%d to determine "
               "wavelength ranges", irecout);
  char *hname = cpl_sprintf("ESO PRO REC%d PARAM%d VALUE", irecout, iparout1),
       *hvalue = cpl_sprintf("%g", lminout),
       *hcomment = cpl_sprintf("set with %s!", idstring);
  cpl_propertylist_update_string(cube->header, hname, hvalue);
  cpl_propertylist_set_comment(cube->header, hname, hcomment);
  cpl_free(hname);
  cpl_free(hvalue);
  hname = cpl_sprintf("ESO PRO REC%d PARAM%d VALUE", irecout, iparout2),
  hvalue = cpl_sprintf("%g", lmaxout),
  cpl_propertylist_update_string(cube->header, hname, hvalue);
  cpl_propertylist_set_comment(cube->header, hname, hcomment);
  cpl_free(hname);
  cpl_free(hvalue);
  cpl_free(hcomment);

  /* save output, then check for errors, and end */
  int rc = muse_datacube_save(cube, oname);
  muse_datacube_delete(cube);
  if (rc == CPL_ERROR_NONE) {
    cpl_msg_info(idstring, "Saved cube to \"%s\".", oname);
  } else {
    cpl_msg_error(idstring, "Error while saving cube to \"%s\": %s (%d)", oname,
                  cpl_error_get_message(), rc);
    rc = 20;
  }

  if (!cpl_errorstate_is_equal(state)) {
    cpl_errorstate_dump(state, CPL_FALSE, muse_cplerrorstate_dump_some);
    rc = 50;
  }
  cpl_memory_dump();
  cpl_end();
  return rc;
}

/**@}*/
