#include <string>
#include <iostream>
#include <iomanip>
#include <vector>
#include <cmath>

#include "TRandom3.h"

#include "JFit/JGradient.hh"
#include "JTools/JElement.hh"
#include "JTools/JQuantile.hh"
#include "JMath/JGauss.hh"
#include "JMath/JConstants.hh"

#include "Jeep/JPrint.hh"
#include "Jeep/JParser.hh"
#include "Jeep/JMessage.hh"

namespace {

  using namespace JPP;


  /**
   * Fit object.
   */
  JGauss fit;
  

  /**
   * Data.
   */
  typedef JElement2D<double, double>     element_type;
  typedef std::vector<element_type>      buffer_type;
  
  buffer_type buffer;


  /**
   * Fit function.
   *
   * \param  gauss              gauss
   * \param  point              point
   * \return                    chi2
   */
  inline double g1(const JGauss& gauss, const element_type& point) 
  {
    const double u  = (point.getX() - gauss.mean) / gauss.sigma;
    const double fs =  gauss.signal * exp(-0.5*u*u) / (sqrt(2.0*PI) * gauss.sigma);
    const double fb =  gauss.background;
    const double f1 =  fs + fb;

    const double y  = point.getY();
    const double dy = (f1 - y) * (f1 - y);

    if (y > 0.0)
      return dy / y;
    else
      return f1;
  }


  /**
   * Get chi2.
   *
   * \param  option            option
   * \return                   chi2
   */
  double getChi2(const int option)
  {
    double chi2 = 0.0;
    
    for (buffer_type::const_iterator i = buffer.begin(); i != buffer.end(); ++i) {
      chi2 += g1(fit, *i);
    }

    return chi2;
  }


  /**
   * Model specific fit parameter.
   */
  struct JGaussEditor :
    public JParameter_t
  {
    /**
     * Constructor.
     *
     * \param  g1                 fit object
     * \param  g2                 step
     */
    JGaussEditor(JGauss& g1, const JGauss& g2) :
      g1(g1),
      g2(g2)
    {}
      
    
    /**
     * Apply step.
     *
     * \param  step              step
     */
    virtual void apply(const double step) override
    {
      g1.add(g2 * step);
    }

    JGauss& g1;
    JGauss  g2;
  };
}


/**
 * \file
 *
 * Program to test JFIT::JGradient algorithm.
 * \author mdejong
 */
int main(int argc, char **argv)
{
  using namespace std;
  using namespace JPP;

  int      numberOfEvents;
  JGauss   gauss;
  JGauss   precision;
  UInt_t   seed;
  int      debug;

  try {

    JParser<> zap("Program to test JGradient algorithm.");

    zap['n'] = make_field(numberOfEvents)      = 0;
    zap['@'] = make_field(gauss)               = JGauss(0.0,  1.0,  1000.0, 10.0);
    zap['e'] = make_field(precision)           = JGauss(0.05, 0.05,   25.0,  25.0);
    zap['S'] = make_field(seed)                = 0;
    zap['d'] = make_field(debug)               = 3;

    zap(argc, argv);
  }
  catch(const exception& error) {
    FATAL(error.what() << endl);
  }

  gRandom->SetSeed(seed);

  if (numberOfEvents == 0) {

    buffer.push_back(element_type(-3.00000,     16.00000));
    buffer.push_back(element_type(-2.50000,     22.00000));
    buffer.push_back(element_type(-2.00000,     70.00000));
    buffer.push_back(element_type(-1.50000,    152.00000));
    buffer.push_back(element_type(-1.00000,    261.00000));
    buffer.push_back(element_type(-0.50000,    337.00000));
    buffer.push_back(element_type( 0.00000,    378.00000));
    buffer.push_back(element_type( 0.50000,    360.00000));
    buffer.push_back(element_type( 1.00000,    253.00000));
    buffer.push_back(element_type( 1.50000,    129.00000));
    buffer.push_back(element_type( 2.00000,     72.00000));
    buffer.push_back(element_type( 2.50000,     22.00000));
    buffer.push_back(element_type( 3.00000,     11.00000));

    JGradient gradient(1000000, 0, 1.0e-4, debug);

    // start values
    
    fit = JGauss(0.5, 0.5, 700.0, 0.0);

    gradient.push_back(JModifier_t("mean",       new JGaussEditor(fit, JGauss(1.0, 0.0, 0.0, 0.0)), 1.0e-2));
    gradient.push_back(JModifier_t("sigma",      new JGaussEditor(fit, JGauss(0.0, 1.0, 0.0, 0.0)), 1.0e-2));
    gradient.push_back(JModifier_t("signal",     new JGaussEditor(fit, JGauss(0.0, 0.0, 1.0, 0.0)), 5.0e-0));
    gradient.push_back(JModifier_t("background", new JGaussEditor(fit, JGauss(0.0, 0.0, 0.0, 1.0)), 5.0e-1));
    
    gradient(getChi2);

  } else {

    JQuantile Q[] = { JQuantile("mean      "),
		      JQuantile("sigma     "),
		      JQuantile("signal    "),
		      JQuantile("background") };

    const size_t nx   = 21;
    const double xmin = -5.0;
    const double xmax = +5.0;
  
    for (int i = 0; i != numberOfEvents; ++i) {
    
      STATUS("event: " << setw(10) << i << '\r'); DEBUG(endl);

      buffer.clear();
    
      for (double x = xmin, dx = (xmax - xmin) / (nx - 1); x < xmax + 0.5*dx; x += dx) {

	const double value = gauss(x);
      
	buffer.push_back(element_type(x, gRandom->Poisson(value)));
      }

      JGradient gradient(1000000, 0, 1.0e-4, debug);

      // start values
    
      fit = JGauss(0.5, 0.5, 700.0, 0.0);

      gradient.push_back(JModifier_t("mean",       new JGaussEditor(fit, JGauss(1.0, 0.0, 0.0, 0.0)), 1.0e-2));
      gradient.push_back(JModifier_t("sigma",      new JGaussEditor(fit, JGauss(0.0, 1.0, 0.0, 0.0)), 1.0e-2));
      gradient.push_back(JModifier_t("signal",     new JGaussEditor(fit, JGauss(0.0, 0.0, 1.0, 0.0)), 5.0e-0));
      gradient.push_back(JModifier_t("background", new JGaussEditor(fit, JGauss(0.0, 0.0, 0.0, 1.0)), 5.0e-1));

      const double chi2 = gradient(getChi2);

      DEBUG("Final value " << fit << endl);
      DEBUG("Chi2 " << chi2 << endl);

      Q[0].put(fit.mean       - gauss.mean);
      Q[1].put(fit.sigma      - gauss.sigma);
      Q[2].put(fit.signal     - gauss.signal);
      Q[3].put(fit.background - gauss.background);
    }

    for (int i = 0; i != sizeof(Q)/sizeof(Q[0]); ++i) {
      NOTICE((i == 0 ? longprint : shortprint) << Q[i]);
    }

    for (int i = 0; i != sizeof(Q)/sizeof(Q[0]); ++i) {
      ASSERT(fabs(Q[i].getMean()) < Q[i].getSTDev());
    }

    ASSERT(Q[0].getSTDev() < precision.mean);
    ASSERT(Q[1].getSTDev() < precision.sigma);
    ASSERT(Q[2].getSTDev() < precision.signal);
    ASSERT(Q[3].getSTDev() < precision.background);
  }

  return 0;
}