/**
  \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.

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

  int band0, band;
  std::string tele0, tele;
  std::vector<std::string> keys = {"BAND", "TELESCOP"};
  hdu.keyWord(keys[0]).value(band0);
  hdu.keyWord(keys[1]).value(tele0);

  if (cli.getOption('e', &files)) {
    for (auto &file: files) {
      std::unique_ptr<CCfits::FITS>
        ferr{new CCfits::FITS(file, CCfits::Read, false, keys)};
      auto &phdu = ferr->pHDU();
      phdu.keyWord(keys[0]).value(band);
      phdu.keyWord(keys[1]).value(tele);
      if (band == band0 && tele == tele0) {
        phdu.read(err);
        MU_EXCEPTION_IF(err.size() != data.size(), fn,
          "Error map '%s' size mismatch (%ud pixels, %ud in parent)",
          file, err.size(), data.size());

        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 *= relerr; }
  MU_INFO(fn, "No error map found for %ld micron band, assuming %g%% of input",
          band0, 100*relerr);
  return err;
}


/**
  \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();

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

  // 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);
    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));

    // Get image data and keywords.
    int band;
    std::string facility;
    bandData.push_back(mapDataType());
    hd.read(bandData.back());
    hd.readAllKeys();
    MU_INFO(fn, "%s (%s): %ldx%ld @ %3ld microns (%s)", file,
            hd.keyWord("OBJECT").value(name),
            hd.axis(0), hd.axis(1), hd.keyWord("BAND").value(band),
            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));
  }

  // 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;
  solve(outMap, outHDU, outFITS.get(), bandData, bandError, bandHDU, cli);

  // Populate header keywords and write data.
  for (unsigned i=0, n=outMap.size(); i != n; i++) {
    setKeys(*outHDU.at(i), *bandHDU[0]);
    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++) {
      int band;
      std::string bunit = "BUNIT", units,
                  comm = bandHDU[i]->getComments(),
                  hist = bandHDU[i]->getHistory();
      bandHDU[i]->readKey(bunit, units);
      bandHDU[i]->keyWord("BAND").value(band);
      std::string name = "BAND" + std::to_string(band);

      std::vector<long> axes = {bandHDU[i]->axis(0), bandHDU[i]->axis(1)};
      auto eHDU = outFITS->addImage(name, mapDataCode, axes);
      eHDU->addKey(bunit, units, "Image flux units");
      eHDU->copyAllKeys(bandHDU[i]);
      eHDU->writeComment(comm);
      eHDU->writeHistory(hist);
      eHDU->write(1, bandData[i].size(), bandData[i]);

      name.insert(0, "d");
      eHDU = outFITS->addImage(name, mapDataCode, axes);
      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->writeHistory(hist);
      eHDU->write(1, bandError[i].size(), bandError[i]);
    }
  }
}


/**
  \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.
*/
void setKeys(CCfits::ExtHDU &hdu, const CCfits::PHDU &hdu0)
{
  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
  auto name = hdu.name().c_str();
  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') {
    hdu.addKey<std::string>("BUNIT", "chi2/DOF", "Goodness of fit units");
  }

  // Copy select keywords from fiducial image.
  auto &inKeys = hdu0.keyWord();
  for (auto &key: {"OBJECT",
                   "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);
    }
  }
}

}  // namespace manticore