//#ifndef __CINT__
#include "RooGlobalFunc.h"
//#endif
#include "RooRealVar.h"
#include "RooGenericPdf.h"
#include "RooArgList.h"
#include "RooFormulaVar.h"
#include "RooDataSet.h"
#include "RooGaussian.h"
#include "RooChebychev.h"
#include "RooExponential.h"
#include "RooBifurGauss.h"
#include "RooAddModel.h"
#include "RooProdPdf.h"
#include "TCanvas.h"
#include "RooPlot.h"
#include "RooHist.h"
#include "RooCBShape.h"
#include "RooPolynomial.h"
#include "RooBinning.h"
#include "TH1.h"
#include "TH2.h"
#include "RooAddPdf.h"
#include "RooProdPdf.h"
#include "RooFitResult.h"
#include "RooGenericPdf.h"
#include "RooLandau.h"
#include "TChain.h"
#include "TCanvas.h"
#include "TAxis.h"
#include "TStyle.h"
#include "TH1F.h"
#include "TAxis.h"
#include "TLine.h"
#include "RooPlot.h"
#include "RooCategory.h"
#include "RooSuperCategory.h"
#include "RooSimultaneous.h"
#include "RooNLLVar.h"
#include "TFile.h"
#include "RooFit.h"
#include "RooGaussModel.h"
#include "RooHistPdf.h"
#include "RooPolyVar.h"

#include "myRooGaussian.cpp"

using namespace RooFit ;
using namespace std;

class RooSelfNormalizedGaussian : public RooAbsPdf {
public:
  RooSelfNormalizedGaussian(const char *name, const char *title,
                            RooAbsReal& _x, RooAbsReal& _mean, RooAbsReal& _sigma) :
  RooAbsPdf(name,title),
  x("x","Observable",this,_x),
  mean("mean","Mean",this,_mean),
  sigma("sigma","Width",this,_sigma)
{}

  RooSelfNormalizedGaussian(const RooSelfNormalizedGaussian& other, const char* name=0) :
    RooAbsPdf(other,name), x("x",this,other.x), mean("mean",this,other.mean),
    sigma("sigma",this,other.sigma) {}
  
  virtual TObject* clone(const char* newname) const override {
    return new RooSelfNormalizedGaussian(*this,newname);
  }

  Bool_t selfNormalized() const override { return true; }
   
  Int_t getAnalyticalIntegral(RooArgSet& allVars, RooArgSet& analVars, const char* rangeName=0) const override {
    if (matchArgs(allVars,analVars,x)) return 1 ;
    if (matchArgs(allVars,analVars,mean)) return 2 ;
    return 0 ;
  }                             

  Double_t analyticalIntegral(Int_t code, const char* rangeName=0) const override {
    assert(code==1 || code==2);
    return 1.0;
  }

protected:

  Double_t evaluate() const override {
    const double arg = x - mean;
    const double sig = sigma;
    return 1./(sig * std::sqrt(TMath::TwoPi())) * std::exp(-0.5*arg*arg/(sig*sig));
  }

private:
  RooRealProxy x ;
  RooRealProxy mean ;
  RooRealProxy sigma ;
};

void unbinned_Kpipi0_RS_2D_fit_updated()  
{

  // get the data in a TTree format
  TFile *f = new TFile("DtoKpipi0_RS_mc14ri_a_2Dfit.root");
  TTree *tree = (TTree*)f->Get("DstD0PiKPiPi0RS_withcut");
  
  // check we have the tree
  cout << "Entries = " << tree->GetEntries() << endl;

  //create both fit observables and get the data in RooFit land
  RooRealVar m("pi0_m_prefit","#it{m}(#it{#pi}^{0})", 0.1, 0.18, "GeV/#it{c}^{2}");
  RooRealVar errm("pi0_errm","#it{m}(#it{#pi}^{0}) error", 0, 0.015, "GeV/#it{c}^{2}");
  RooDataSet* dataxy = new RooDataSet("data","data",tree, RooArgSet(m, errm));

  
  //*  S i n g l e    G a u s s i o n     R e s o l u t i o n    f u n c t i o n
  //  ------------------------------------------------------------------------
  
  errm.setRange("SB1",0,0.0041);
  errm.setRange("SB2",0.0041,0.0055);
  errm.setRange("SB3",0.0055,0.015);

  std::vector<RooPolyVar> mus;

  // mean: Create function f(y) = a0 + a1*y (errm range: 0,0.0041)
  RooRealVar p00("p00", "p00", 1.34010e-01, -1, 1); // 1.34010e-01  
  RooRealVar p01("p01", "p01", -2.72535e-01, -1, 1); // -2.72535e-01
  mus.emplace_back("mu0", "mean0", errm, RooArgSet(p00, p01));

  // mean: Create function f(y) = a0 + a1*y + a2*y^2 (errm range: 0.0041,0.0055)
  RooRealVar p10("p10", "p10", 9.22697e-02, -0.2, 0.2); // 9.22697e-02  
  RooRealVar p11("p11", "p11", 1.68039e+01, -1, 1); // 1.68039e+01
  RooRealVar p12("p12", "p12", -1.68719e+03, -0.1, 0.1); // -1.68719e+03
  mus.emplace_back("mu1", "mean1", errm, RooArgSet(p10, p11, p12));

  // mean: Create function f(y) = a0 + a1*y + a2*y^2 (errm range: 0.0055,0.015)
  RooRealVar p20("p20", "p20", 1.31621e-01, -1, 1); // 1.31621e-01  
  RooRealVar p21("p21", "p21", 3.81109e-01, -1, 1); // 3.81109e-01
  mus.emplace_back("mu2", "mean2", errm, RooArgSet(p20, p21));

  RooFormulaVar mu("mu", "pi0_errm < 0.0041 ? mu0 : (pi0_errm < 0.0055 ? mu1 : mu2)", RooArgList(errm, mus[0], mus[1], mus[2]));

  // to set the parameters for `mu` in a given range constant in the fit
  auto setConstant = [&](RooAbsReal const& real, bool constant) {
      auto parameters = real.getParameters(*dataxy);
      for (auto param : *parameters) {
          if(auto realVar = dynamic_cast<RooRealVar*>(param)) {
            realVar->setConstant(constant);
          }
      }
      delete parameters;
  };

  //Sigma: Create function f(y) = 1 + a/x + b/x^2 
  RooRealVar a("a", "a", -1.54699e-03); // -1.54699e-03  
  RooRealVar b("b", "b", 1.15198e-05);  // 1.15198e-05
  RooFormulaVar sigma("sigma", "@0+@1+@2/@0", RooArgList(errm,a,b));

  //myRooGaussian gauss("gauss", "1st gaussian PDF", m, mu, sigma);
  //RooGaussian gauss("gauss", "1st gaussian PDF", m, mu, sigma);
  RooSelfNormalizedGaussian gauss("gauss", "1st gaussian PDF", m, mu, sigma);


  // C o n s t r u c t  p d f   f o r   e r r o r   m a s s
  // -----------------------------------------------------------------
  
  // Construct a histogram pdf to describe the shape of the ErrM distribution
  RooDataSet datamerr("datamerr","datamerr", tree, RooArgSet(errm));

  RooDataHist *expHistDterr = datamerr.binnedClone();
  RooHistPdf pdfErr("pdfErr", "pdfErr", errm, *expHistDterr);
 

  // C o n s t r u c t   c o n d i t i o n a l   p r o d u c t   d e c a y 
   // -----------------------------------------------------------------------------
  
  // Construct production of conditional decay_dm(dt|dterr) with empirical pdfErr(dterr)
  RooProdPdf pdf("pdf", "pdf(m|sigma_m)*pdf(sigma_m)", pdfErr, RooFit::Conditional(gauss, m));

  
  
  // perform the fit in the three separat ranges, each time with only the parameters relevant in a given range floating
  setConstant(mus[1], true);
  setConstant(mus[2], true);
  pdf.fitTo(*dataxy, Save(true), Strategy(2), Range("SB1"), PrintLevel(-3))->Print();
  setConstant(mus[0], true);
  setConstant(mus[1], false);
  pdf.fitTo(*dataxy, Save(true), Strategy(2), Range("SB2"), PrintLevel(-3))->Print();
  setConstant(mus[1], true);
  setConstant(mus[2], false);
  pdf.fitTo(*dataxy, Save(true), Strategy(2), Range("SB3"), PrintLevel(-3))->Print();
  setConstant(mus[0], false);
  setConstant(mus[1], false);

  // reset the fitrange attribute before plotting, so the plot doesn't think we have only fitted a subrange
  pdf.setStringAttribute("fitrange", nullptr);
  
 /*-----------------------------------PLOT THE RESULTS-------------------------------------*/

  //-----------------------Canvas and pad
  TCanvas *canvas = new TCanvas("canvas","canvas", 700, 800);  //  700, 800
 
  Double_t xlow, ylow, xup, yup;
  canvas->GetPad(0)->GetPadPar(xlow,ylow,xup,yup);
  canvas->Divide(1,2);

  TVirtualPad *upPad = canvas->GetPad(1);
  upPad->SetPad(xlow,ylow+0.25*(yup-ylow),xup,yup);
  
  TVirtualPad *dwPad = canvas->GetPad(2);
  dwPad->SetPad(xlow,ylow,xup,ylow+0.25*(yup-ylow));

  //------------------------Plot pi0_M_prefit projections
  RooPlot *xframe = m.frame();
  dataxy->plotOn(xframe);
  // const double nData = dataxy->sumEntries("", "SB2");
  // cout<<" nentries in data are = "<<nData<<endl;
  pdf.plotOn(xframe, ProjWData(*dataxy));
  pdf.paramOn(xframe);
  canvas->cd(1);
  //pdf.plotOn(xframe,Components(gauss),LineStyle(kDashed),LineColor(kRed), ProjWData(*dataxy));
  xframe->Draw();

  cout << "chi^2 = " << xframe->chiSquare() << endl;
  return;

  //----------------------- pull distribution of pi0_M_prefit
  RooHist* hpull = xframe->pullHist();
  hpull->SetFillStyle(1001);
  hpull->SetFillColor(1);
  for(int i=0;i<hpull->GetN();++i) hpull->SetPointError(i,0.0,0.0,0.0,0.0);
  RooPlot* pullplot = m.frame(Title(" "));
  pullplot->addPlotable(hpull,"B");
  pullplot->SetYTitle("Pull mass");
  pullplot->SetMinimum(-4.);
  pullplot->SetMaximum(4.);
  pullplot->GetXaxis()->SetLabelSize(0.1);
  pullplot->GetYaxis()->SetLabelSize(0.07);
  canvas->cd(2);
  pullplot->Draw();

  //------------------------- adding reference line in pull distribution
  double xmin1 = 0.1; 
  double xmax1 = 0.18;
  TLine *line = new TLine(xmin1,0.0,xmax1,0.0);
  TLine *line1 = new TLine(xmin1,3.0,xmax1,3.0);
  TLine *line2 = new TLine(xmin1,-3.0,xmax1,-3.0);
 
  line->SetLineColor(kRed); 
  line->SetLineWidth(3); 
  line1->SetLineColor(kRed);
  line2->SetLineColor(kRed);
  line1->SetLineStyle(2);
  line2->SetLineStyle(2);
  line->Draw("SAME"); 
  line1->Draw("SAME");
  line2->Draw("SAME");
  
  

  //------------------------- Plot pi0_ErrM projections
  RooPlot *yframe = errm.frame();
  dataxy->plotOn(yframe);
  pdf.plotOn(yframe);

  canvas->Update();
  

  // TCanvas *c = new TCanvas("c", "c", 700, 800);
  // //canvas->cd(2);
  // gPad->SetLeftMargin(0.15);
  // yframe->GetYaxis()->SetTitleOffset(1.6);
  // yframe->Draw();
 

 
}