#include <TH1.h>
#include <TCanvas.h>
#include <RooRealVar.h>
#include <RooConstVar.h>
#include <RooDataHist.h>
#include <RooCFunction3Binding.h>
#include <RooCFunction4Binding.h>
#include <RooFormulaVar.h>
#include <RooFFTConvPdf.h>
#include <RooExtendPdf.h>
#include <RooAddPdf.h>
#include <RooArgusBG.h>
#include <RooSimultaneous.h>
#include <RooFitResult.h>
#include <RooPlot.h>
#include <RooNumIntConfig.h>
#include <RooMsgService.h>

#include <cmath>

using namespace RooFit;
using namespace std;

static TString Mean() { return "m_{#phi}"; }
static TString Sigma() { return "#sigma"; }
static TString Gamma() { return "#Gamma"; }


inline double GetKMass() { return 493.677; }
inline double GetKThreshold() { return 2 * GetKMass(); }

TString GetPolyParName(const int i) { return TString::Format("a%d", i); }


inline
double
qGaussian(const double x, const double mean, const double sigma, const double q)
{
  const double power = 1.0 / (1.0 - q);
  const double alpha = 0.5 * (x - mean) * (x - mean) / sigma / sigma;
  return std::pow(1 + (q - 1) * alpha, power);
}


double
NA49RelativisticBW(const double m, const double mean, const double gamma)
{
  static const double mk = GetKMass();
  static const double th = GetKThreshold();

  if (m <= th)
    return 0;

  const double mk2 = mk * mk;
  const double m2 = m * m;
  const double m02 = mean * mean;

  const double q2 = m2 / 4.0 - mk2;
  const double q02 = m02 / 4.0 - mk2;
  const double q = std::sqrt(q2);
  const double q0 = std::sqrt(q02);

  const double gammaOfM = 2 * gamma * std::pow(q / q0, 3) * q02 / (q2 + q02);

  const double denominator = (m2 - m02) * (m2 - m02) +
      m02 * gammaOfM * gammaOfM;

  return m * gammaOfM / denominator;
}


inline
double
Shift(const double x, const double xmax, const double threshold)
{
  return xmax + threshold - x;
}


void
TuneIntegrator()
{
  // Change default RooFit numeric integrator to a faster one
  RooNumIntConfig* config = RooAbsReal::defaultIntegratorConfig();
  config->method1D().setLabel("RooAdaptiveGaussKronrodIntegrator1D");

  auto& section =
    config->getConfigSection("RooAdaptiveGaussKronrodIntegrator1D");
  section.setRealValue("maxSeg", 50);
  section.setCatLabel("method", "31Points");

  // Prevent numeric integrator from polluting cout with error messages
  RooMsgService::instance().getStream(0).removeTopic(RooFit::Integration);
}


void
test_chi2()
{
  const int nbins = 104;
  const double xmin = 986;
  const double xmax = 1090;
  TH1D h("h", "", nbins, xmin, xmax);
  const double contents[nbins] = {0, 2590, 7358, 9975, 12081, 13687, 15329, 16454, 18108, 19239, 20359, 21073, 22134, 23223, 23943, 24876, 25793, 26098, 27256, 27923, 28734, 29422, 30179, 30643, 31692, 32043, 32915, 33858, 34714, 36359, 37819, 40216, 42516, 43786, 43178, 41529, 40518, 39831, 38854, 39058, 38993, 39022, 39445, 39505, 39512, 39959, 40214, 40576, 40917, 40711, 41102, 41481, 42059, 41911, 42198, 42646, 42614, 42816, 42734, 43536, 44221, 43904, 43717, 44536, 44636, 44828, 44906, 44764, 45636, 45238, 45847, 45734, 45878, 46046, 46431, 46670, 46859, 46599, 47105, 47512, 46868, 47659, 47322, 47711, 47822, 47947, 48347, 47997, 48128, 48667, 48320, 48457, 48921, 48805, 49126, 48959, 49412, 48748, 49250, 49184, 49365, 49342, 49755, 50120};
  for (int i = 1; i <= nbins; ++i)
    h.SetBinContent(i, contents[i - 1]);

  RooRealVar x("x", "x", xmin, xmax);
  x.setRange("norm", xmin, xmax);
  x.setRange("cache", xmin, xmax);
  x.setBins(10000, "cache");

  RooDataHist data("data", "data", x, &h);

  RooRealVar nsig("Nsig", "Nsig", 62958.53, 0, 2e5);

  RooRealVar width(Gamma(), Gamma(), 4.266, "MeV/c^{2}");
  RooRealVar mean(Mean(), Mean(), 1019.461, 1015.0, 1025.0, "MeV/c^{2}");
  RooRealVar sigma(Sigma(), Sigma(), 1.0, 0.05, 2.5, "MeV/c^{2}");
  RooRealVar q("q", "q", 1.5, "");

  auto& det = *bindPdf("det", qGaussian, x, RooConst(0), sigma, q);
  auto& bw = *bindPdf("bw", NA49RelativisticBW, x, mean, width);
  RooFFTConvPdf signal("signal", "signal", x, bw, det);
  RooExtendPdf esig("esig", "esig", signal, nsig, "norm");

  auto& mX = *bindFunction("mX", Shift, x, RooConst(xmax), RooConst(GetKThreshold()));

  RooRealVar k(GetPolyParName(0), GetPolyParName(0), -1.0, -10.0, 10.0);
  RooRealVar p(GetPolyParName(1), GetPolyParName(1), 0.5, 0.0, 1.0);
  RooArgusBG bkg("bkg", "bkg", mX, RooConst(xmax), k, p);

  RooRealVar nbkg("nbkg", "nbkg", 3921654.47, 0, 5e6);
  RooExtendPdf ebkg("ebkg", "ebkg", bkg, nbkg);

  RooAddPdf model("model", "model", RooArgList(ebkg, esig));


  TCanvas* can = new TCanvas("test", "");
  can->SaveAs(".pdf[");
  RooFitResult* result = nullptr;
  RooPlot* frame = nullptr;

  // With Offset() the fit fails with default integrator, uncomment the below
  // line to make it converge again.
  // TuneIntegrator();
  result = model.fitTo(data, PrintLevel(1), Save(), Minimizer("Minuit2", "Minimize"), Offset());
  result->Print("V");
  frame = x.frame();
  data.plotOn(frame);
  model.plotOn(frame,  Normalization(1, RooAbsReal::RelativeExpected));
  frame->SetTitle("fitTo");
  frame->Draw();
  can->SaveAs(".pdf");

  result = model.chi2FitTo(data, PrintLevel(-1), Save(), Minimizer("Minuit2", "Minimize"), Extended(true));
  result->Print("V");
  frame = x.frame();
  data.plotOn(frame);
  model.plotOn(frame,  Normalization(1, RooAbsReal::RelativeExpected));
  frame->SetTitle("chi2FitTo");
  frame->Draw();

  can->SaveAs(".pdf");

  can->SaveAs(".pdf]");
}