/**
  \file

  \brief SED fitting routines.

  \warning
  This file assumes GSL version 2.1 or earlier.
  With v2.2, the non-linear least-squares API was completely rewritten!!

  \who \kpr
*/
#include <ctime>
#include <thread>

#include <gsl/gsl_multifit_nlin.h>

#include "Graybody.h"
#include "PACS.h"
#include "SPIRE.h"

namespace mu = mutils;

/// Package namespace.
namespace manticore {

/// Observation fitting data structure.
struct GrayData {
  /// Constructor.
  GrayData(Graybody *g, std::vector<mapDataType> *o, size_t n) :
    gray(g), outMap(o)
  {
    Fobs.resize(n);
    Ferr.resize(n);
    band.resize(n);
  }

  Graybody *gray;              ///< Graybody emission model.

  std::vector<long> axes;      ///< Image dimensions.

  std::vector<double> Fobs,    ///< Observed flux in each band.
                      Ferr,    ///< Flux uncertainty in each band.
                      dnu0;    ///< Effective bandwidth.

  std::vector<const Detector*>
                      band;    ///< Band detectors.

  std::vector<mapDataType>
                     *outMap;  ///< Output maps.
};

/**
  \brief Observation fitting function.
  \param[in]  x Model parameters {T, Sigma}.
  \param[in]  p Fitting data structure.
  \param[out] f Output functions.
  \return GSL_SUCCESS.
*/
int f_gray(const gsl_vector *x, void *p, gsl_vector *f)
{
  GrayData *gd = static_cast<GrayData *>(p);
  const auto &Fobs = gd->Fobs, &Ferr = gd->Ferr;
  const auto &gray = *gd->gray;

  size_t n = Ferr.size();
  double T = gsl_vector_get(x, 0), S = gsl_vector_get(x, 1);
  for (size_t i=0; i != n; i++) {
    auto Di = gd->band[i];
    double iFerr = 1.0/Ferr[i];
    gsl_vector_set(f, i, (gray.F(T, S, Di) - Fobs[i])*iFerr);
  }

  return GSL_SUCCESS;
}

/**
  \brief Observation fitting function Jacobian.
  \param[in]  x Model parameters {T, Sigma}.
  \param[in]  p Fitting data structure.
  \param[out] J Output Jacobian.
  \return GSL_SUCCESS.
*/
int df_gray(const gsl_vector *x, void *p, gsl_matrix *J)
{
  GrayData *gd = static_cast<GrayData *>(p);
  const auto &Ferr = gd->Ferr;
  const auto &gray = *gd->gray;

  size_t n = Ferr.size();
  double T = gsl_vector_get(x, 0), S = gsl_vector_get(x, 1);
  for (size_t i=0; i != n; i++) {
    auto Di = gd->band[i];
    double iFerr = 1.0/Ferr[i];
    gsl_matrix_set(J, i, 0, gray.dF_dT(T, S, Di)*iFerr);
    gsl_matrix_set(J, i, 1, gray.dF_dS(T, S, Di)*iFerr);
  }

  return GSL_SUCCESS;
}

/**
  \brief Observation fitting function plus Jacobian.
  \param[in]  x Model parameters {T, Sigma}.
  \param[in]  p Fitting data structure.
  \param[out] f Output functions.
  \param[out] J Output Jacobian.
  \return GSL_SUCCESS.
*/
int fdf_gray(const gsl_vector *x, void *p, gsl_vector *f, gsl_matrix *J)
{ return f_gray(x, p, f), df_gray(x, p, J); }


/// Current time a la std::asctime().
std::string asctime()
{
  char buf[64];
  struct tm tmbuf;
  auto t = std::time(nullptr);
  std::string rtn = asctime_r(localtime_r(&t, &tmbuf), buf);
  rtn.pop_back();  // Remove newline.
  return rtn;
}


/**
  \brief Stage 1 solver.

  Single-plane mass/temperature solution.

  \param[out] outMap     Output result data.
  \param[in]  bandData   Input band data.
  \param[in]  bandError  Input band (1-sigma) uncertainties.
  \param[in]  cli        Command line options and arguments.
  \param[in]  gdata0     Model parameters.
  \param[in]  id         Thread ID number (0-based).
  \param[in]  nThreads   Number of threads.
*/
void solveStage1(
                 const std::vector<mapDataType> &bandData,
                 const std::vector<mapDataType> &bandError,
                 const mu::CommandLine &cli,
                 const GrayData &gdata0, unsigned id, unsigned nThreads)
{
  const char *const fn = "solveStage1";
  size_t nBands = bandData.size(), nParam = 2,
         minBands = cli.getInteger('n', 3);
  const double imu_g = 1.0/(2.75*m_H);
  auto &Tmap  = gdata0.outMap->at(0),
       &dTmap = gdata0.outMap->at(1),
       &Nmap  = gdata0.outMap->at(2),
       &dNmap = gdata0.outMap->at(3),
       &Cmap  = gdata0.outMap->at(4);

  // Thread-local storage.
  Graybody gray{*gdata0.gray};
  GrayData gdata{gdata0};
  gdata.gray = &gray;

  // GSL v1.x least-squares fitting setup.
  gsl_multifit_fdfsolver *solver =
    gsl_multifit_fdfsolver_alloc(gsl_multifit_fdfsolver_lmsder, // or lmder
                                 nBands, nParam);
  //
  gsl_multifit_function_fdf fdf;
  fdf.f = f_gray;
  fdf.df = df_gray;
  fdf.fdf = fdf_gray;
  fdf.n = nBands;
  fdf.p = nParam;
  fdf.params = &gdata;
  //
  // Initial guess.
  gsl_vector *x0 = gsl_vector_alloc(nParam);
  gsl_vector_set(x0, 0, 15.0);
  gsl_vector_set(x0, 1,  1.0);
  //
  // Covariance storage.
  gsl_matrix *cov = gsl_matrix_alloc(nParam, nParam);

  // Loop over map pixels.
  const double relerr = cli.getDouble('A', 1e-2);
  size_t maxIter = std::max(1L, cli.getInteger('M', 25)),
         px = bandData[0].size(),
         i0 = (id*px)/nThreads, i1 = ((id+1)*px)/nThreads;
  for (unsigned i=i0; i != i1; i++) {
    // Set target data.
    unsigned nOk = 0;
    for (size_t j=0; j != nBands; j++) {
      gdata.Fobs[j] = gdata.dnu0[j]*bandData[j][i];
      gdata.Ferr[j] = gdata.dnu0[j]*bandError[j][i];

      if (gdata.Fobs[j] > 0.0 && gdata.Ferr[j] > 0.0) {
        nOk++;
      } else {
        gdata.Fobs[j] = 0.0;
        gdata.Ferr[j] = 1e10;  // Weight-out bad data.
      }
    }

    // Progress display.
    MU_INFO_IF(id == 0 && i%(i1/10) == 0, fn,
               "Fitting solution %2.0f%% complete @ %s",
               (100.0*i)/i1, asctime());

    // Skip invalid pixels.
    if (nOk < minBands) {
      Tmap[i] = dTmap[i] = Nmap[i] = dNmap[i] = Cmap[i] = NAN;
      continue;
    }

    // Initialize solver.
    gsl_multifit_fdfsolver_set(solver, &fdf, x0);

    // Iterate on solution.
    size_t iter = 0;
    do {
      gsl_multifit_fdfsolver_iterate(solver);
      MU_DEBUG(fn, "[%2lu] T = %.2f  S = %.4f", iter,
               gsl_vector_get(solver->x, 0), gsl_vector_get(solver->x, 1));
    } while (gsl_multifit_test_delta(solver->dx, solver->x, 0.0, relerr) ==
             GSL_CONTINUE && ++iter != maxIter);

    // Results.
    gsl_multifit_covar(solver->J, 1e-8, cov);
    double chi2 = 0.0;
    for (size_t k=0; k != nBands; k++) {
      auto fk = gsl_vector_get(solver->f, k);
      chi2 += fk*fk;
    }
    //
    double chi2_dof = chi2/(nBands-nParam), vscale = std::max(1.0, chi2_dof),
           T = gsl_vector_get(solver->x, 0), S = gsl_vector_get(solver->x, 1),
           dT = sqrt(vscale * gsl_matrix_get(cov, 0, 0)),
           dS = sqrt(vscale * gsl_matrix_get(cov, 1, 1));
    MU_DEBUG(fn, "T = %.2f(%.2f), S = %.4f(%.4f), chi2/DOF = %.3f\n",
             T, dT, S, dS, chi2_dof);
    //
    Tmap[i] = T,       dTmap[i] = dT;
    Nmap[i] = S*imu_g, dNmap[i] = dS*imu_g;
    Cmap[i] = chi2_dof;
    //
    MU_WARN_IF(iter == maxIter, fn,
               "Fitting failed for pixel (%ld,%ld): dT/T = %.4g, dN/N = %.4g",
               1+i%gdata.axes[0], 1+i/gdata.axes[0],
               gsl_vector_get(solver->dx, 0)/T,
               gsl_vector_get(solver->dx, 1)/S);

    // Chain next initial guess to previous solution.
    gsl_vector_set(x0, 0, T);
    gsl_vector_set(x0, 1, S);
  }

  // Free storage.
  gsl_multifit_fdfsolver_free(solver);
  gsl_vector_free(x0);
  gsl_matrix_free(cov);
}


/**
  \brief Performs least-squares fitting of band SEDs.

  Given input band continuum fluxes and uncertainties, calculates the best-fit
  temperature and column density maps (including uncertainties and reduced
  chi-square information). This routine does not hard-wire a particular number
  of temperature/density components; an appropriate fitting engine is selected
  based on the `--components NCOMP` option (default: 1).
  Currently only `NCOMP` = 1 is supported.

  For `NCOMP` = 1 the output FITS file will contain the following named image
  extensions:
  - **T**      (gas temperature)
  - **dT**     (1-sigma uncertainty in **T**)
  - **N(H2)**  (N(H2) gas column density)
  - **dN(H2)** (1-sigma uncertainty in **N**)
  - **Chi2**   (fit chi-squared per degree of freedom)

  \param[out]    outMap     Output result data.
  \param[in,out] outHDU     Output FITS headers.
  \param[in,out] outFITS    Output FITS file.
  \param[in]     bandData   Input band data.
  \param[in]     bandError  Input band (1-sigma) uncertainties.
  \param[in]     bandHDU    Input band FITS headers.
  \param[in]     cli        Command line options and arguments.
*/
void solve(std::vector<mapDataType> &outMap,
           std::vector<CCfits::ExtHDU *> &outHDU,
           CCfits::FITS *outFITS,
           const std::vector<mapDataType> &bandData,
           const std::vector<mapDataType> &bandError,
           const std::vector<CCfits::PHDU*> &bandHDU,
           const mu::CommandLine &cli)
{
  const char *const fn = "solve";

  // Dust model.
  bool extend = !cli.getOption('p');
  Dust dust(cli.getString('d', "OH5"), cli.getDouble('g', 100.0));
  MU_INFO(fn, "Using %s source detector response",
          extend ? "extended" : "point");
  MU_INFO(fn, "Dust model: %s, gas-to-dust ratio = %g",
          dust.name(), dust.rho());

  // Allocate result maps.
  const double relerr = cli.getDouble('A', 1e-2);
  std::vector<long> axes = {bandHDU[0]->axis(0), bandHDU[0]->axis(1)};
  auto px = axes[0]*axes[1];
  char comm[1024];
  snprintf(comm, sizeof(comm),
           "%s sources, %s dust, gas/dust = %g, relerr = %.1e",
           (extend ? "Extended" : "Point"), dust.name().c_str(),
           dust.rho(), relerr);
  //
  for (auto name: {"T", "dT", "N(H2)", "dN(H2)", "Chi2/DOF"}) {
    outMap.push_back(mapDataType(0.0, px));
    outHDU.push_back(outFITS->addImage(name, mapDataCode, axes));
    outHDU.back()->writeDate();
    outHDU.back()->writeComment(comm);
  }

  // Band detector setup.
  // Currently supports PACS/SPIRE only!
  size_t nBands = bandData.size();
  std::vector<std::unique_ptr<Detector>> detect;
  for (size_t j=0; j != nBands; j++) {
    int band;
    bandHDU[j]->keyWord("BAND").value(band);

    auto d = (band < 200 ? (Detector *)new PACS(band, extend)
                         : (Detector *)new SPIRE(band, extend));
    d->init();
    detect.emplace_back(d);
  }

  // Emission model.
  constexpr double MJy = 1e-17;  // MJy -> erg/cm^2/s/Hz
  Graybody gray{dust, 0.1*relerr};
  GrayData gdata{&gray, &outMap, nBands};
  for (size_t j=0; j != nBands; j++) {
    gdata.band[j] = detect[j].get();
    gdata.dnu0.push_back(detect[j]->freqWidth() * MJy);
  }
  gdata.axes = axes;

  // Create threads and process maps.
  unsigned nThreads = std::max(1L, cli.getInteger('t', 1));
  std::vector<std::thread> threads;
  for (unsigned i=1; i != nThreads; i++) {
    threads.emplace_back(&solveStage1, bandData, bandError, cli,
                         gdata, i, nThreads);
  }
  solveStage1(bandData, bandError, cli, gdata, 0, nThreads);
  //
  for (auto &t: threads) { t.join(); }

  MU_INFO(fn, "Solution complete @ %s", asctime());
}

}  // namespace manticore