/**
  \file

  \brief Map processing infrastructure.

  \who \kpr
*/
#include "manticore.h"
#include "version.h"

namespace mu = mutils;

/// Package namespace.
namespace manticore {

/**
  \brief Generates image error map.

  If any `--error-map` options were supplied and an entry matching
  the **BAND** and **TELESCOP** keywords of the input is found, the error map
  is read from the corresponding image. Otherwise, the value of `--error-pct`
  (default: 3.0) is multiplied by *data* to create the error map.

  \param[in] data  Input image map data.
  \param[in] hdu   Input image map FITS header.
  \param[in] cli   Command line options and arguments.
  \param[in] ll    Cutout lower-left  vertex.
  \param[in] ur    Cutout upper-right vertex.
  \param[in] st    Cutout stride (cutout disabled if st[0] = 0).

  \return Image 1-sigma error map.
*/
mapDataType getError(const mapDataType &data, CCfits::PHDU &hdu,
                     const mu::CommandLine &cli,
                     const std::vector<long> &ll,
                     const std::vector<long> &ur,
                     const std::vector<long> &st)
{
  const char *const fn = "getError";
  mapDataType err;
  mu::CommandLine::Arguments files;

  int band0 = getBand(hdu), band;
  std::string tele0, tele, bandKey = "BAND", teleKey = "TELESCOP";
  hdu.keyWord(teleKey).value(tele0);

  if (cli.getOption('e', &files)) {
    for (auto &file: files) {
      std::unique_ptr<CCfits::FITS>
        ferr{new CCfits::FITS(file, CCfits::Read, false, {bandKey, teleKey})};
      auto &phdu = ferr->pHDU();
      band = getBand(phdu);
      phdu.keyWord(teleKey).value(tele);
      if (band == band0 && tele == tele0) {
        if (st[0]) { phdu.read(err, ll, ur, st); }
        else       { phdu.read(err); }
        MU_EXCEPTION_IF(phdu.axis(0) != hdu.axis(0) ||
                        phdu.axis(1) != hdu.axis(1), fn,
          "Error map '%s' size mismatch (%lux%lu pixels, %lux%lu in parent)",
          file, phdu.axis(0), phdu.axis(1), hdu.axis(0), hdu.axis(1));

        MU_INFO(fn, "Using error map '%s' for %ld micron band", file, band0);
        return err;
      }
    }
  }

  double relerr = 0.01*cli.getDouble('E', 3.0);
  err = data;
  for (auto &elem: err) { elem = fabs(elem * relerr); }
  MU_INFO(fn, "No error map found for %ld micron band, assuming %g%% of input",
          band0, 100*relerr);
  return err;
}

/**
  \brief Reset cutout header coordinates.

  \param[in,out] hdu  Input image map FITS header.
  \param[in]     ll   Lower-left  cutout vertex (cutout disabled if zero).
*/
void resetRefPix(CCfits::HDU &hdu, const std::vector<long> &ll)
{
  if (ll[0]) {
    double pix;
    hdu.keyWord("CRPIX1").setValue(1+hdu.keyWord("CRPIX1").value(pix)-ll[0]);
    hdu.keyWord("CRPIX2").setValue(1+hdu.keyWord("CRPIX2").value(pix)-ll[1]);
  }
}

/**
  \brief Processes multi-wavelength maps into temperature/column density maps.

  This routine handles the higher-level management required to generate mass
  and temperature images from the input continuum data, including:
  - reading the input FITS images and headers
  - reading or generating input map uncertainties
  - performing input compatibility checks
  - populating the output FITS header and data structures
  - writing the output FITS file to disk

  \param[in] cli  Command line options and arguments.
*/
void process(const mu::CommandLine &cli)
{
  const char *const fn = "process";
  auto &files = cli.arguments();
  bool modelMap = cli.getOption('M');

  // Sanity check.
  MU_EXCEPTION_IF(files.size() < 2, fn,
                  "Minimum of 2 input bands required (%lu given); try --help",
                  files.size());

  // Cutout processing.
  std::vector<long> ll{0,0}, ur{10000,10000}, st{0,0};
  auto cutout = cli.getString('C', mutils::nullStr);
  if (cutout != mutils::nullStr) {
    unsigned nc, nr, r0, c0;
    MU_EXCEPTION_IF(4 != sscanf(cutout.c_str(), "%ux%u@%u,%u",
                                &nc, &nr, &c0, &r0), fn,
                    "Invalid syntax '--cutout %s'; try --help", cutout);
    MU_INFO(fn, "Processing %lux%lu cutout at (%lu,%lu)",
            nc, nr, c0, r0);
    ll[0] = c0,      ll[1] = r0;
    ur[0] = c0+nc-1, ur[1] = r0+nr-1;
    st[0] = st[1] = 1;
  }

  // Read the FITS images.
  // It is assumed ALL images have the same pixel scale and dimensions,
  // even if the beams are mismatched.
  std::vector<std::unique_ptr<CCfits::FITS>> img;
  std::vector<CCfits::PHDU*> bandHDU;
  std::vector<mapDataType> bandData, bandError;
  for (auto &file: files) {
    img.emplace_back(new CCfits::FITS(file, CCfits::Read, true));
    bandHDU.push_back(&img.back()->pHDU());

    // Axis check.
    std::string name;
    auto &hd0 = *bandHDU.front(), &hd = *bandHDU.back();
    int naxes = hd.axes();
    MU_EXCEPTION_IF(naxes != 2, fn, "%s has %ld axes (expected 2)", file,naxes);

    // Get image data and keywords.
    std::string facility;
    bandData.push_back(mapDataType());
    if (st[0]) { hd.read(bandData.back(), ll, ur, st); }
    else       { hd.read(bandData.back()); }
    MU_EXCEPTION_IF(modelMap &&
                    bandData.back().size() < modelMapLen*modelMapLen, fn,
                    "Input maps too small for %ldx%ld model",
                    modelMapLen, modelMapLen);
    MU_ERROR_IF(hd.axis(0) != hd0.axis(0) || hd.axis(1) != hd0.axis(1), fn,
                    "%s dimensions are %ldx%ld (fiducial is %ldx%ld)",
                    file, hd.axis(0), hd.axis(1), hd0.axis(0), hd0.axis(1));

    hd.readAllKeys();
    MU_INFO(fn, "%s (%s): %ldx%ld @ %3ld microns (%s)", file,
            hd.keyWord("OBJECT").value(name),
            hd.axis(0), hd.axis(1), getBand(hd),
            hd.keyWord("TELESCOP").value(facility));

    // Unit check.
    std::string units;
    hd.readKey("BUNIT", units);
    MU_EXCEPTION_IF(units != "MJy/sr", fn,
                    "%s map units are '%s' (need MJy/sr)", file, units);
  }

  // Find or create associated error maps.
  for (unsigned i=0; i != bandData.size(); i++) {
    bandError.push_back(getError(bandData[i], *bandHDU[i], cli, ll, ur, st));
  }

  //
  // Output maps (one image extension per map).
  auto mtfile = cli.getString('o', "mtcore.fits");
  if (cli.getOption('F')) { mtfile.insert(0, 1, '!'); }
  std::unique_ptr<CCfits::FITS>
    outFITS(new CCfits::FITS(mtfile, BYTE_IMG, 0, nullptr));

  // Fit the mass and temperature maps.
  std::vector<mapDataType> outMap;
  std::vector<CCfits::ExtHDU*> outHDU;
  std::vector<long> axis{
    modelMap ? modelMapLen : ll[0] ? 1+ur[0]-ll[0] : bandHDU[0]->axis(0),
    modelMap ? modelMapLen : ll[0] ? 1+ur[1]-ll[1] : bandHDU[0]->axis(1)
  };
  solve(outMap, outHDU, outFITS.get(), bandData, bandError, bandHDU, axis,
        cli, ll, ur, st);

  // Populate header keywords and write data.
  for (unsigned i=0, n=outMap.size(); i != n; i++) {
    setKeys(*outHDU.at(i), *bandHDU[0], cli, ll);
    outHDU[i]->write(1, outMap[i].size(), outMap[i]);
  }

  if (!cli.getOption('B')) {
    // Append input band data to FITS output.
    for (unsigned i=0, n=bandData.size(); i != n; i++) {
      std::string name = "BAND" + std::to_string(getBand(*bandHDU[i]));
      std::string bunit = "BUNIT", units,
                  comm = bandHDU[i]->getComments();
      bandHDU[i]->readKey(bunit, units);

      if (modelMap) { axis = {modelMapLen, modelMapLen}; }

      auto eHDU = outFITS->addImage(name, mapDataCode, axis);
      eHDU->addKey(bunit, units, "Image flux units");
      eHDU->copyAllKeys(bandHDU[i]);
      eHDU->writeComment(comm);
      eHDU->write(1, axis[0]*axis[1], bandData[i]);
      resetRefPix(*eHDU, ll);

      name.insert(0, "d");
      eHDU = outFITS->addImage(name, mapDataCode, axis);
      eHDU->writeComment("ERROR map used for the associated parent image");
      eHDU->addKey(bunit, units, "Error flux units");
      eHDU->copyAllKeys(bandHDU[i]);
      eHDU->writeComment(comm);
      eHDU->write(1, axis[0]*axis[1], bandError[i]);
      resetRefPix(*eHDU, ll);
    }
  }
}


/**
  \brief Sets header keywords for an output map.

  Combines information from the input FITS headers and the command line
  arguments to populate keywords in the output FITS file.

  \param[in,out] hdu   Output FITS header.
  \param[in]     hdu0  Fiducial input image header.
  \param[in]     cli   Command line options and arguments.
  \param[in]     ll    Lower-left cutout vertex (ignored if zero).
*/
void setKeys(CCfits::ExtHDU &hdu, const CCfits::PHDU &hdu0,
             const mu::CommandLine &cli, const std::vector<long> &ll)
{
  const char *const fn = "setKeys";

  // Local keywords.
  hdu.addKey("CREATOR", std::string("Manticore v") + version,
             "FITS file creator");
  //
  // Units are determined from the extension name:
  // - N/dN: 1/cm^2
  // - T/dT: K
  // - Chi2: chi2/DOF
  // - nIter: counts
  // - diff*: absolute/relative flux difference
  // - frac*: unity
  auto name = hdu.name();
  if (name[0] == 'N' || name[1] == 'N') {
    hdu.addKey<std::string>("BUNIT", "cm^-2", "Gas column density units");
  }
  if (name[0] == 'T' || name[1] == 'T') {
    hdu.addKey<std::string>("BUNIT", "K", "Gas temperature units");
  }
  if (name[0] == 'C' || ! name.compare(0, 4, "dChi")) {
    hdu.addKey<std::string>("BUNIT", "chi2/DOF", "Goodness of fit units");
  }
  if (name[0] == 'n') {
    hdu.addKey<std::string>("BUNIT", "counts", "Fitter iterations");
  }
  if (! name.compare(0, 4, "diff")) {
    if (cli.getOption('a')) {
      hdu.addKey<std::string>("BUNIT", "MJy/sr",
                              "Absolute differential flux units");
    } else {
      hdu.addKey<std::string>("BUNIT", "sigma",
                              "Relative differential flux units");
    }
  }
  if (! name.compare(0, 4, "frac")) {
    hdu.addKey<std::string>("BUNIT", "ratio", "Cold flux fraction");
  }

  // Copy select keywords from fiducial image.
  auto &inKeys = hdu0.keyWord();
  for (auto &key: {"OBJECT", "RADESYS",
                   "CTYPE1", "CRVAL1", "CRPIX1", "CDELT1",
                   "CTYPE2", "CRVAL2", "CRPIX2", "CDELT2",
                  }) {
    auto kp = inKeys.find(key);
    if (kp != inKeys.end()) {
      hdu.addKey(kp->second);
    } else {
      MU_EXCEPTION(fn, "Fiducial image missing keyword '%s'", key);
    }
  }

  // Sky coordinates.
  resetRefPix(hdu, ll);
}

}  // namespace manticore