/* BEGIN software license
 *
 * MsXpertSuite - mass spectrometry software suite
 * -----------------------------------------------
 * Copyright 2009--2026 by 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
 */


#pragma once


////////////////////////////// Stdlib includes


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


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


/////////////////////// pappsomspp includes
#include <pappsomspp/core/trace/trace.h>


/////////////////////// Local includes
#include "MsXpS/export-import-config.h"

/*
Bitwise stuff (from StackOverflow)

It is sometimes worth using an enum to name the bits:

enum ThingFlags = {
ThingMask  = 0x0000,
ThingFlag0 = 1 << 0,
ThingFlag1 = 1 << 1,
ThingError = 1 << 8,
}

Then use the names later on. I.e. write

thingstate |= ThingFlag1;
thingstate &= ~ThingFlag0;
if (thing & ThingError) {...}

to set, clear and test. This way you hide the magic numbers from the rest of
your code.
 */

namespace MsXpS
{
namespace libXpertMassCore
{
namespace Math
{


class DECLSPEC Gaussian
{
  public:
  struct LevMarqGaussianFunctor
  {
    enum
    {
      InputsAtCompileTime = Eigen::Dynamic,
      ValuesAtCompileTime = Eigen::Dynamic
    };

    typedef Eigen::Matrix<double, InputsAtCompileTime, 1> InputType;
    typedef Eigen::Matrix<double, ValuesAtCompileTime, 1> ValueType;
    typedef Eigen::Matrix<double, ValuesAtCompileTime, InputsAtCompileTime>
      JacobianType;

    // BEST PRACTICE: Eigen::Ref accepts VectorXd, Map, or Block WITHOUT copying
    const Eigen::Ref<const Eigen::VectorXd> x_data;
    const Eigen::Ref<const Eigen::VectorXd> y_data;
    const Eigen::Ref<const Eigen::VectorXd> w;
    const int num_points;
    const double fixed_sigma;

    // Constructor now accepts x, y, and weights
    LevMarqGaussianFunctor(const Eigen::Ref<const Eigen::VectorXd> &x,
                           const Eigen::Ref<const Eigen::VectorXd> &y,
                           const Eigen::Ref<const Eigen::VectorXd> &weights,
                           double sigma)
      : x_data(x),
        y_data(y),
        w(weights),
        num_points(x.size()),
        fixed_sigma(sigma)
    {
    }

    int
    inputs() const
    {
      return 3;
    } // A, mu, C

    int
    values() const
    {
      return num_points;
    }

    int
    operator()(const InputType &params, ValueType &fvec) const
    {
      double A     = params[0];
      double mu    = params[1];
      double C     = params[2];
      double sigma = fixed_sigma;

      if(sigma <= 1e-6)
        sigma = 1e-6; // Prevent division by zero

      for(int i = 0; i < num_points; ++i)
        {
          double x     = x_data[i];
          double y     = y_data[i];
          double arg   = (x - mu) / sigma;
          double model = A * std::exp(-0.5 * arg * arg) + C;

          // Apply weight to residual.
          // NOTE: Standard Weighted Least Squares minimizes sum(w * r^2),
          // which requires multiplying the residual by std::sqrt(w).
          // (If your old code incorrectly used just w[i], change
          // std::sqrt(w[i]) to w[i])
          double weight = std::sqrt(w[i]);
          fvec[i]       = weight * (model - y);
        }
      return 0;
    }

    int
    df(const InputType &params, JacobianType &fjac) const
    {
      double A     = params[0];
      double mu    = params[1];
      double sigma = fixed_sigma;

      if(sigma <= 1e-6)
        sigma = 1e-6;

      fjac.resize(num_points, 3);

      for(int i = 0; i < num_points; ++i)
        {
          double x        = x_data[i];
          double arg      = (x - mu) / sigma;
          double exp_term = std::exp(-0.5 * arg * arg);

          // Apply the SAME weight to the Jacobian derivatives
          double weight = std::sqrt(w[i]);

          fjac(i, 0) = weight * exp_term;                     // d/dA
          fjac(i, 1) = weight * A * exp_term * (arg / sigma); // d/dmu
          fjac(i, 2) = weight * 1.0;                          // d/dC
        }
      return 0;
    }
  };

  explicit Gaussian();

  virtual ~Gaussian();

  static std::tuple<double, double, double>
  linearLeastSquaresFit(const pappso::Trace &input_trace,
                        pappso::Trace &fit_trace);

  static std::tuple<double, double, double, double>
  levenbergMarquardtFit(const pappso::Trace &input_trace,
                        pappso::Trace &fit_trace);
};

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