/* BEGIN software license
 *
 * MsXpertSuite - mass spectrometry software suite
 * -----------------------------------------------
 * Copyright (C) 2009--2024 Filippo Rusconi
 *
 * http://www.msxpertsuite.org
 *
 * This file is part of the MsXpertSuite project.
 *
 * The MsXpertSuite project is the successor of the massXpert project. This
 * project now includes various independent modules:
 *
 * - massXpert, model polymer chemistries and simulate mass spectrometric data;
 * - mineXpert, a powerful TIC chromatogram/mass spectrum viewer/miner;
 *
 * 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 3 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, see <http://www.gnu.org/licenses/>.
 *
 * END software license
 */


////////////////////////////// Stdlib includes
// #include <math.h>
#include <cstdio>
#include <cmath>
#include <iostream>
#include <iomanip>


////////////////////////////// Qt includes
#include <QDebug>


/////////////////////// Eigen3 includes
#include <Eigen/Dense>
#include <unsupported/Eigen/NonLinearOptimization>


////////////////////////////// Local includes
#include "MsXpS/libXpertMassCore/Math/Generic.hpp"
#include "MsXpS/libXpertMassCore/Math/Gaussian.hpp"
#include "MsXpS/libXpertMassCore/MassPeakShaper.hpp"

namespace MsXpS
{
namespace libXpertMassCore
{
namespace Math
{


/*!
\class MsXpS::libXpertMassCore::Math::Gaussian
\inmodule libXpertMassCore
\ingroup XpertMassCoreMath
\inheaderfile MsXpS/libXpertMassCore/Math/Gaussian.hpp

\brief The Gaussian class provides a number of utility features for Gaussian
fitting.
*/

/*!
\brief Performs a Gaussian fit of the \a input_trace and sets the computed fit
data to \a fit_trace.

The Gaussian fit is computed using the least squares method.

The returned tuple has these members:

\list

\li \e {Member 0}: Gaussian mu (µ), the abscissa of the curve's apex
\li \e {Member 1}: A (amplitude), the intensity at the apex
\li \e {Member 2}: FWHM (the full width at half maximum)

\endlist

If an error occurs, or if input_trace is empty or no fit could be computed, a
tuple will all three members set to 0 is returned.
*/
std::tuple<double, double, double>
Gaussian::linearLeastSquaresFit(const pappso::Trace &input_trace,
                                pappso::Trace &fit_trace)
{

  // At this point we can fit a Gaussian on the trace.

  int i, j, k, n, N;

  std::cout.precision(4); // set precision
  std::cout.setf(std::ios::fixed);

  N = input_trace.size();

  if(!N)
    return {0, 0, 0};

  // qDebug() << "The number of local maxima: " << N;

  double *x     = new double[N];
  double *y     = new double[N];
  double *log_y = new double[N];

  pappso::DataPoint most_intense_datapoint;

  for(size_t iter = 0; iter < input_trace.size(); ++iter)
    {
      pappso::DataPoint data_point = input_trace.at(iter);

      if(data_point.y > most_intense_datapoint.y)
        most_intense_datapoint = data_point;

      x[iter] = data_point.x;
      y[iter] = data_point.y;

      // Transform the y values into ln(y) so that we transform something that
      // might be exponential to something that is polynomial (Caruana stuff).

      // std::log is the natural log (ln) in the stdlib, contrary to std::log10.

      log_y[iter] = std::log(data_point.y);
    }

  // Set the polynomial degree to 2 because we want a Gaussian fit and that is
  // what will provide the parameters for the Gaussian fit.
  n = 2;

  double *X = new double[2 * n + 1]; // Array that will store the values of
  // sigma(xi),sigma(xi^2),sigma(xi^3)....sigma(xi^2n)
  for(i = 0; i < 2 * n + 1; i++)
    {
      X[i] = 0;
      for(j = 0; j < N; j++)
        X[i] =
          X[i] + pow(x[j], i); // consecutive positions of the array will store
      // N,sigma(xi),sigma(xi^2),sigma(xi^3)....sigma(xi^2n)
    }

  // B is the Normal matrix(augmented) that will store the
  // equations. So it is an array of arrays.
  double **B = new double *[n + 1];

  for(int iter = 0; iter < n + 1; ++iter)
    B[iter] = new double[n + 2];

  // a is for value of the final coefficient
  double *a = new double[n + 1];

  // Build the Normal matrix by storing the
  // corresponding coefficients at the right positions
  // except the last column of the matrix
  for(i = 0; i <= n; i++)
    for(j = 0; j <= n; j++)
      B[i][j] = X[i + j];

  // Array to store the values of
  // sigma(yi),sigma(xi*yi),sigma(xi^2*yi)...sigma(xi^n*yi)
  double *Y = new double[n + 1];
  for(i = 0; i < n + 1; i++)
    {
      Y[i] = 0;
      for(j = 0; j < N; j++)
        // Y[i]=Y[i]+pow(x[j],i)*y[j];        //consecutive positions will store
        // sigma(yi),sigma(xi*yi),sigma(xi^2*yi)...sigma(xi^n*yi)
        Y[i] =
          Y[i] + pow(x[j], i) * log_y[j]; // consecutive positions will store
      // sigma(yi),sigma(xi*yi),sigma(xi^2*yi)...sigma(xi^n*yi)
    }
  for(i = 0; i <= n; i++)
    B[i][n + 1] = Y[i]; // load the values of Y as the last column of B(Normal
  // Matrix but augmented)
  n = n + 1; // n is made n+1 because the Gaussian Elimination part below was
  // for n equations, but here n is the degree of polynomial and for
  // n degree we get n+1 equations

  std::cout << "\nThe Normal(Augmented Matrix) is as follows:\n";
  for(i = 0; i < n; i++) // print the Normal-augmented matrix
    {
      for(j = 0; j <= n; j++)
        std::cout << B[i][j] << std::setw(16);
      std::cout << "\n";
    }

  for(i = 0; i < n;
      i++) // From now Gaussian Elimination starts(can be ignored) to solve
    // the set of linear equations (Pivotisation)
    for(k = i + 1; k < n; k++)
      if(B[i][i] < B[k][i])
        for(j = 0; j <= n; j++)
          {
            double temp = B[i][j];
            B[i][j]     = B[k][j];
            B[k][j]     = temp;
          }

  for(i = 0; i < n - 1; i++) // loop to perform the gauss elimination
    for(k = i + 1; k < n; k++)
      {
        double t = B[k][i] / B[i][i];
        for(j = 0; j <= n; j++)
          B[k][j] =
            B[k][j] - t * B[i][j]; // make the elements below the pivot elements
                                   // equal to zero or elimnate the variables
      }
  for(i = n - 1; i >= 0; i--) // back-substitution
    { // x is an array whose values correspond to the values of x,y,z..
      a[i] = B[i][n]; // make the variable to be calculated equal to the rhs of
      // the last equation
      for(j = 0; j < n; j++)
        if(j != i) // then subtract all the lhs values except the coefficient of
          // the variable whose value is being calculated
          a[i] = a[i] - B[i][j] * a[j];
      a[i] = a[i] / B[i][i]; // now finally divide the rhs by the coefficient of
                             // the variable to be calculated
    }

  // std::cout << "\nThe values of the coefficients are as follows:\n";
  // for(i = 0; i < n; i++)
  // std::cout << "x^" << i << "=" << a[i]
  //<< std::endl; // Print the values of x^0,x^1,x^2,x^3,....
  // std::cout << "\nHence the fitted Polynomial is given by:\ny=";
  // for(i = 0; i < n; i++)
  // std::cout << " + (" << a[i] << ")"
  //<< "x^" << i;
  // std::cout << "\n";

  double a_coeff        = a[0];
  double b_coeff        = a[1];
  double b_coeff_square = b_coeff * b_coeff;
  double c_coeff        = a[2];

  // std::cout << "a_coeff: " << a_coeff << "\nb_coeff: " << b_coeff
  //<< " b_coeff_square: " << b_coeff_square << "\nc_coeff: " << c_coeff
  //<< std::endl;

  double gaussian_mu = -b_coeff / (2 * c_coeff);
  // double gaussian_mu_square = gaussian_mu * gaussian_mu;

  double gaussian_sigma = sqrt(-1 / (2 * c_coeff));
  double fwhm           = 2 * gaussian_sigma;

  // std::cout << "gaussian_mu: " << gaussian_mu
  //<< "\ngaussian_sigma: " << gaussian_sigma
  //<< " gaussian_sigma_square: " << gaussian_sigma_square << std::endl;

  // Now perform the computation over the x values that were entered
  // initially.

  // From the Hongwei Guo article

  double A = std::exp(a_coeff - (b_coeff_square / (4 * c_coeff)));

  // std::cout
  //<< "Calculating A according to: A = exp(a_coeff - (b_coeff_square / (4 "
  //"* c_coeff)))"
  //<< std::endl;

  // std::cout << "A is now: " << A << std::endl;

  delete[] x;
  delete[] y;
  delete[] X;
  delete[] Y;

  delete[] a;

  // qDebug() << "n has value:" << n;

  for(int iter = 0; iter < n; ++iter)
    {
      // qDebug() << "Iterating at iter:" << iter;
      delete[] B[iter];
      // qDebug() << "deleted B at iter:" << iter;
    }

  // At this point use our facility to compute the Gaussian trace.

  // This was the old version.
  // MassPeakShaperConfig mass_peak_shaper_config(
  // 0, fwhm, "H", 1, 500, 0, 1,
  // Enums::MassPeakShapeType::GAUSSIAN);

  MassPeakShaperConfig mass_peak_shaper_config;
  mass_peak_shaper_config.setResolution(0);
  mass_peak_shaper_config.setFwhm(fwhm);
  mass_peak_shaper_config.setPointCount(200);
  mass_peak_shaper_config.setReferencePeakMz(most_intense_datapoint.x);
  mass_peak_shaper_config.setMzStep(0);
  mass_peak_shaper_config.setMassPeakShapeType(
    Enums::MassPeakShapeType::GAUSSIAN);

  MassPeakShaper mass_peak_shaper(gaussian_mu, A, mass_peak_shaper_config);

  mass_peak_shaper.computePeakShape();
  fit_trace = std::move(mass_peak_shaper.getTrace());

  return std::tuple<double, double, double>{gaussian_mu, A, fwhm};
}

/*!
\brief Performs a Gaussian fit of the \a input_trace and sets the computed fit
data to \a fit_trace.

The Gaussian fit is computed using the Levenberg-Marquardt method. The
Levenberg-Marquardt functor that is used for the minimization is fed with
parameters obtained from running Math::Generic::weightedCenterOfMass.

The returned tuple has four members:

\list

\li \e {Member 0}: Gaussian mu (µ), the abscissa of the curve's apex
\li \e {Member 1}: A (amplitude), the intensity at the apex
\li \e {Member 2}: sigma [related to the FWHM according to:  FWHM = 2 * sqrt(2 *
std::log(2)) * sigma]
\li \e {Member 3}: b, the baseline value

\endlist

If an error occurs, or if input_trace is empty or no fit could be computed, a
tuple will all three members set to 0 is returned.

\sa Math::Gaussian::linearLeastSquaresFit, Math::Generic::weightedCenterOfMass
*/
std::tuple<double, double, double, double>
Gaussian::levenbergMarquardtFit(const pappso::Trace &input_trace,
                                pappso::Trace &fit_trace)
{

  // First we need to get starting data by using the weighted center of mass
  // results:

  std::tuple<double, double, double, double> weighted_center_of_mass_results =
    Math::Generic::weightedCenterOfMass(input_trace);

  double mu            = std::get<0>(weighted_center_of_mass_results);
  double A             = std::get<1>(weighted_center_of_mass_results);
  // sigma is half of 95% of the curve area
  double sigma         = std::get<2>(weighted_center_of_mass_results);
  double b             = std::get<3>(weighted_center_of_mass_results);

  qDebug().noquote() << qSetRealNumberPrecision(10)
                     << "Weighted center of mass results:";

  qDebug().noquote() << "avg_mz:" << mu << "max_intensity:" << A
                     << "sigma:" << sigma << "b0:" << b;

  Eigen::VectorXd x_data =
    Eigen::VectorXd::Map(input_trace.xValues().data(), input_trace.size());
  Eigen::VectorXd y_data =
    Eigen::VectorXd::Map(input_trace.yValues().data(), input_trace.size());

  Eigen::VectorXd w = y_data.cwiseMax(1e-6).cwiseInverse(); // 1/y weights

  Eigen::VectorXd params(3);
  params << A, mu, b;

  LevMarqGaussianFunctor functor(x_data, y_data, w, sigma);

  Eigen::LevenbergMarquardt<LevMarqGaussianFunctor> lm(functor);

  lm.parameters.maxfev = 50; // safety cap; will converge in ~4

  // Perform the minimization.
  lm.minimize(params);

  // Retrieve the values.
  A  = params[0];
  mu = params[1];
  b  = params[2];

  qDebug() << qSetRealNumberPrecision(5) << "A:" << A << "mu:" << mu
           << "b:" << b;

  // At this point use our facility to compute the Gaussian trace.

  MassPeakShaperConfig mass_peak_shaper_config;
  mass_peak_shaper_config.setResolution(0);
  mass_peak_shaper_config.setFwhm(2 * sqrt(2 * std::log(2)) * sigma);
  mass_peak_shaper_config.setPointCount(200);
  mass_peak_shaper_config.setReferencePeakMz(mu);
  mass_peak_shaper_config.setMzStep(0);
  mass_peak_shaper_config.setMassPeakShapeType(
    Enums::MassPeakShapeType::GAUSSIAN);

  MassPeakShaper mass_peak_shaper(mu, A, mass_peak_shaper_config);

  mass_peak_shaper.computePeakShape();
  fit_trace = std::move(mass_peak_shaper.getTrace());

  return {mu, A, sigma, b};
}

} // namespace Math
} // namespace libXpertMassCore
} // namespace MsXpS
