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

#include "TROOT.h"
#include "TFile.h"
#include "TH1D.h"
#include "TRandom3.h"
#include "TMath.h"

#include "JFit/JSimplex.hh"
#include "JTools/JArray.hh"
#include "JTools/JArrayIterator.hh"
#include "JTools/JGrid.hh"
#include "JTools/JQuantile.hh"
#include "JROOT/JRootToolkit.hh"

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

namespace {

  using JTOOLS::JArray;
  using TMath::PoissonI;

  static const int NUMBER_OF_DIMENSIONS  =  3;


  /**
   * Type definition of a point in ND-space.
   */
  typedef JArray<NUMBER_OF_DIMENSIONS, double>  point_type;


  /**
   * Model.
   */
  struct model_type :
    public JArray<1 + 2 * NUMBER_OF_DIMENSIONS, double>  
  {
    /**
     * Default constructor.
     */
    model_type() :
      JArray<1 + 2 * NUMBER_OF_DIMENSIONS, double>()
    {}


    /**
     * Constructor.
     *
     * \param  value            first value
     * \param  args             remaining values
     */
    template<class ...Args>
    model_type(argument_type value, const Args& ...args) :
      JArray<1 + 2 * NUMBER_OF_DIMENSIONS, double>(value, args...)
    {}


    /**
     * Model function.
     *
     * \param  point              point
     * \return                    value
     */
    inline double operator()(const point_type& point) const
    {
      double value = (*this)[0];

      for (int i = 0; i != NUMBER_OF_DIMENSIONS;++i) {
	value *= gauss(point[i] - (*this)[2*i + 1], (*this)[2*i + 2]);
      }
      
      return value;
    }


    /**
     * Gauss function (normalised to 1 at x = 0).
     *
     * \param  x                    x
     * \param  sigma                sigma
     * \return                      function value
     */
    static inline double gauss(const double x, const double sigma)
    {
      const double u = x / sigma;

      if (fabs(u) < 10.0)
	return exp(-0.5*u*u);
      else
	return 0.0;
    }
  };


  /**
   * Hit.
   */
  struct hit_type :
    public point_type
  {
    /**
     * Constructor.
     *
     * \param  point              point
     * \param  value              value
     * \param  sigma              sigma
     */
    hit_type(const point_type& point, const double value, const double sigma) :
      point_type(point),
      value(value),
      sigma(sigma)
    {}

    double value;
    double sigma;
  };
 

  /**
   * Fit function based on classical chi2.
   *
   * \param  model              model
   * \param  hit                hit
   * \return                    chi2
   */
  inline double chi2(const model_type& model, const hit_type& hit) 
  {
    const double u = (model(hit) - hit.value) / hit.sigma;

    return u*u;
  }
 

  /**
   * Fit function based on logarithm of likelihood.
   *
   * \param  model              model
   * \param  hit                hit
   * \return                    chi2
   */
  inline double likelihood(const model_type& model, const hit_type& hit) 
  {
    return -log(PoissonI(hit.value, model(hit)));
  }
}


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

  typedef JArray<2, double> array_type;

  string      outputFile;
  int         numberOfEvents;
  array_type  parameters;
  UInt_t      seed;
  int         debug;

  try {

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

    zap['o'] = make_field(outputFile)          = "simplex.root";
    zap['n'] = make_field(numberOfEvents);
    zap['M'] = make_field(parameters, "model parameters") = array_type(100.0, 1.0);
    zap['S'] = make_field(seed)                = 0;
    zap['d'] = make_field(debug)               = 2;

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


  gRandom->SetSeed(seed);


  const double top   = parameters[0];
  const double sigma = parameters[1];

  model_type model;
  
  model[0] = top;

  for (int i = 0; i != NUMBER_OF_DIMENSIONS; ++i) {
    model[i*2 + 1] = gRandom->Uniform(-sigma, +sigma);
    model[i*2 + 2] = gRandom->Gaus(sigma, 0.2*sigma);
  }

  const JGrid<double> grid(11, -3.0*sigma, +3.0*sigma);


  JSimplex<model_type> simplex;

  JSimplex<model_type>::debug               =  debug;
  JSimplex<model_type>::MAXIMUM_ITERATIONS  =  10000;

  double (*fit)(const model_type&, const hit_type&) = likelihood;  // fit function


  vector<TH1D*> H1;

  TH1D h0("chi2/NDF", NULL, 200, 0.0, 10.0);

  H1.push_back(new TH1D(MAKE_CSTRING("p" << 0), NULL, 101, -20.0, +20.0));

  for (size_t i =1; i != model.size(); ++i) {
    H1.push_back(new TH1D(MAKE_CSTRING("p" << i), NULL, 101, -0.1, +0.1));
  }


  for (int counter = 0; counter != numberOfEvents; ++counter) {

    STATUS("event: " << setw(10) << counter << '\r'); DEBUG(endl);


    // generate data
    
    vector<hit_type> data;

    for (JArrayIterator<NUMBER_OF_DIMENSIONS, double> g1(grid); g1; ++g1) {

      const JArray<NUMBER_OF_DIMENSIONS, double> p1 = *g1;

      const double value = gRandom->Poisson(model(p1));
      const double sigma = (value > 1.0 ? sqrt(value) : 0.7);

      data.push_back(hit_type(p1, value, sigma));	  
    }


    // start values

    JQuantile Q[NUMBER_OF_DIMENSIONS];

    double value = 0.0;

    for (const auto& hit : data) {

      for (int i = 0; i != NUMBER_OF_DIMENSIONS; ++i) {
	Q[i].put(hit[i], hit.value);
      }

      if (hit.value > value) {

	simplex.value[0] = hit.value;

	for (int i = 0; i != NUMBER_OF_DIMENSIONS; ++i) {
	  simplex.value[2*i + 1] = hit[i];
	}

	value = hit.value;
      }
    }

    for (int i = 0; i != NUMBER_OF_DIMENSIONS; ++i) {
      simplex.value[2*i + 2] = 1.0 * Q[i].getSTDev();
    }


    // steps

    simplex.step.resize(2*NUMBER_OF_DIMENSIONS + 1);

    model_type step;

    simplex.step[0]    = model_type();
    simplex.step[0][0] = 1.0;

    for (int i = 0; i != NUMBER_OF_DIMENSIONS; ++i) {

      simplex.step[2*i + 1]          = model_type();
      simplex.step[2*i + 1][2*i + 1] = 0.1;

      simplex.step[2*i + 2]          = model_type();
      simplex.step[2*i + 2][2*i + 2] = 0.1;
    }

    DEBUG("start values " << simplex.value << endl);

    for (size_t i = 0; i != simplex.step.size(); ++i) {
      DEBUG("step: " << setw(2) << i << ' ' << simplex.step[i] << endl);
    }


    // fit

    const double chi2 = simplex(fit, data.begin(), data.end());
    const int    ndf  = data.size()  -  simplex.step.size();


    h0.Fill(chi2 / ndf);

    for (size_t i = 0; i != model.size(); ++i) {
      H1[i]->Fill(model[i] - simplex.value[i]);
    }

    DEBUG("chi2 / NDF" << FIXED(12,3) << chi2 << " / " << ndf << endl);

    DEBUG("final values " << simplex.value << endl);
    DEBUG("model values " << model         << endl);

    if (debug >= debug_t) {

      for (const auto& hit : data) {

	cout << "hit: ";

	for (int i = 0; i != NUMBER_OF_DIMENSIONS; ++i) {
	  cout << FIXED(7,3) << hit[i] << ' ';
	}

	const double value = simplex.value(hit);

	cout << FIXED(7,3) << hit.value << ' '
	     << FIXED(7,3) << value  << ' '
	     << FIXED(7,3) << (hit.value - value) / hit.sigma << endl;
      }
    }

  }
  STATUS(endl);
  

  if (outputFile != "") {

    TFile out(outputFile.c_str(), "recreate");

    out << h0;

    for (size_t i = 0; i != sizeof(H1)/sizeof(H1[0]); ++i) {
      out << *(H1[i]);
    }

    out.Write();
    out.Close();
  }
}