#include "RooDataSet.h"
#include "RooRealVar.h"
#include "RooGaussian.h"
#include "RooAddPdf.h"
#include "RooProdPdf.h"
#include "RooAddition.h"
#include "RooProduct.h"
#include "TCanvas.h"
#include "RooChebychev.h"
#include "RooAbsPdf.h"
#include "RooFit.h"
#include "RooFitResult.h"
#include "RooPlot.h"
#include "RooAbsArg.h"
#include "RooWorkspace.h"
#include "RooStats/ProfileLikelihoodCalculator.h"
#include "RooStats/HypoTestResult.h"
#include <string>

// use this order for safety on library loading
using namespace RooFit;
using namespace RooStats;

//____________________________________
void myExample_02(int seed = 2)
{

   gRandom = new TRandom3(0);
   gRandom->SetSeed(seed);

   Double_t lumi = 150;
   Double_t xSecBkgd = 800;
   Double_t xSecSign = 0.5;

   const int nBin = 100;
   const Double_t low = 150;
   const Double_t high = 1150;

   Double_t zmass = 750;
   Double_t zwidth = 25;

   //------------------------//
   // build data and signal  //
   //------------------------//

   // SIGNAL is a gaussian and BCKG is an EXPO
   auto fun_sgnl = new TF1("fun_sgnl", "gaus(0)", low, high);
   auto fun_bkgd = new TF1("fun_bkgd", "expo(0)", low, high);
   fun_sgnl->SetParameters(20, zmass, zwidth);
   fun_bkgd->SetParameters(10, -0.013);

   // number of signal events
   int nSGNL = lumi * xSecSign;

   // number of background events
   int nBKGD = lumi * xSecBkgd;

   // Produce SIGNAL histo
   auto hs = new TH1F("hs1", "", nBin, low, high);
   hs->FillRandom("fun_sgnl", nSGNL); 

   // Produce hbk1 and hbk2 histo for signal and background
   auto hbk1 = new TH1F("hbk11", "", nBin, low, high);
   auto hbk2 = new TH1F("hbk12", "", nBin, low, high);
   hbk1->FillRandom("fun_bkgd", nBKGD);
   hbk2->FillRandom("fun_bkgd", nBKGD);

   // data is the final data set with signal
   // SUM hbk1+hs
   TH1F *data = (TH1F *)hbk1->Clone("data");
   data->SetName("data");
   data->Add(hs);

   // bkgd is the final data set for background (no signal)
   // SUM hbk2+hs2
   TH1F *bkgd = (TH1F *)hbk2->Clone("bkgd");
   bkgd->SetName("bkgd");

   cout << "Entries data - bkgd: " << data->GetEntries() << " - " << bkgd->GetEntries() << endl;

   cout << "//-------------------------------//" << endl;
   cout << "// START DATA ANALYSIS ORIGINAL  //" << endl;
   cout << "//-------------------------------//" << endl;

   //========//
   // STEP 1 //
   //========//

   RooRealVar invMass("invMass", "M_{inv}", low, high, "GeV");
   RooRealVar mu1("mu1", "signal strength in units of SM expectation", 0.001, 0., 1.0);
   RooRealVar expo1("expo1", "expo1", -1.0, -0.00001);
   RooExponential fitExpoFun1("background", "background", invMass, expo1);

   // convert data and bkgd TH1 into a RooDataHist
   RooDataHist hist_data("hist_data", "hist_data", invMass, Import(*data));
   RooDataHist hist_bkgd("hist_bkgd", "hist_bkgd", invMass, Import(*bkgd));

   // FIT Background with an expo function
   auto result = fitExpoFun1.fitTo(hist_bkgd, SumW2Error(true), RooFit::Save(true));
   result->Print();
   expo1.setConstant(kTRUE);

   // Signal is a PDF from the histo used for populating the data histo
   RooDataHist hist_sgnl("hist_sgnl", "hist_sgnl", invMass, Import(*hs));
   RooHistPdf  pdf_sgnl("pdf_sgnl", "pdf_sgnl", invMass, hist_sgnl, 0);

   // full model
   RooAddPdf model1("model1", "sigModel+bkgd shapes", RooArgList(pdf_sgnl, fitExpoFun1), RooArgList(mu1));
   RooArgSet poi1(mu1);
   ProfileLikelihoodCalculator plc1(hist_data, model1, poi1);
   RooArgSet *nullParams1 = (RooArgSet *)poi1.snapshot();
   nullParams1->setRealValue("mu1", 0);
   plc1.SetNullParameters(*nullParams1);

   HypoTestResult *htr1 = plc1.GetHypoTest();
   cout << "************ Significance 1: " << htr1->Significance() << endl;

   // plot signal candidates with background model and components for channel 2
   TCanvas *c1 = new TCanvas();
   RooPlot *xframe1 = invMass.frame();
   hist_data.plotOn(xframe1, DataError(RooAbsData::SumW2));
   model1.plotOn(xframe1);
   model1.plotOn(xframe1, Components(fitExpoFun1), LineStyle(kDashed), LineColor(kRed));
   hist_bkgd.plotOn(xframe1, DataError(RooAbsData::SumW2), MarkerColor(kGreen), LineColor(kGreen));
   xframe1->SetTitle("fit data + signal + background CH 1");
   xframe1->SetMinimum(1e-1);
   xframe1->SetMaximum(5e4);
   xframe1->Draw();
   c1->SetLogy();
   c1->Modified();
   c1->Update();

   //========//
   // STEP 2 //
   //========//

   RooRealVar mu2("mu2", "signal strength in units of SM expectation", 0.001 , 0., 1.0);

   // SIGNAL is a gaussian PDF
   RooRealVar mean1("mean1", "mean of gaussians", low, high);
   RooRealVar sigma1("sigma1", "width of gaussians", zwidth*0.50, zwidth*1.50);
   RooGaussian pdf_sgnl_g("pdf_sgnl_g", "Signal component 1", invMass, mean1, sigma1);
   mean1.setVal(zmass);
   mean1.setConstant(kTRUE);
   sigma1.setVal(zwidth);

   // full model
   RooAddPdf model2("model2", "sigModel+bkgd shapes", RooArgList(pdf_sgnl_g, fitExpoFun1), RooArgList(mu2));
   RooArgSet poi2(mu2);
   ProfileLikelihoodCalculator plc2(hist_data, model2, poi2);
   RooArgSet *nullParams2 = (RooArgSet *)poi2.snapshot();
   nullParams2->setRealValue("mu2", 0);
   plc2.SetNullParameters(*nullParams2);

   HypoTestResult *htr2 = plc2.GetHypoTest();
   cout << "************ Significance 2: " << htr2->Significance() << endl;

   // plot signal candidates with background model and components for channel 2
   TCanvas *c2 = new TCanvas();
   RooPlot *xframe2 = invMass.frame();
   hist_data.plotOn(xframe2, DataError(RooAbsData::SumW2));
   model2.plotOn(xframe2);
   model2.plotOn(xframe2, Components(fitExpoFun1), LineStyle(kDashed), LineColor(kRed));
   hist_bkgd.plotOn(xframe2, DataError(RooAbsData::SumW2), MarkerColor(kGreen), LineColor(kGreen));
   xframe2->SetTitle("fit data + signal + background CH 1");
   xframe2->SetMinimum(1e-1);
   xframe2->SetMaximum(5e4);
   xframe2->Draw();
   c2->SetLogy();
   c2->Modified();
   c2->Update();

   //========//
   // STEP 3 //
   //========//

   Double_t p0width = 1.30053;
   Double_t p1width = 0.00133281;
   Double_t p2width = 4.44654e-05;
   TF1 *widthExp = new TF1("widthExp", "[0]+[1]*x+[2]*x^2", 40, 2060);
   widthExp->SetParameter(0, p0width);
   widthExp->SetParameter(1, p1width);
   widthExp->SetParameter(2, p2width);

   float xVal[nBin];
   float yVal[nBin];
   std::fill(xVal, xVal+nBin, 0);
   std::fill(yVal, yVal+nBin, 0);

   RooRealVar mu3("mu3", "signal strength in units of SM expectation", 0.5, 0., 1);

   // scan for BKGD+SIGNAL and Bkgd only
   int nCount = 0;
   Double_t stepMass = int((high - low) / nBin);
   for (Double_t masScan = low; masScan < high; masScan += stepMass, nCount++)
   {

      xVal[nCount] = masScan;

      Double_t zpWidth = widthExp->Eval(masScan);
      cout << "====================================================" << endl;
      cout << "Start mass point fit: " << nCount << " " << masScan << " " << zpWidth << endl;
      cout << "====================================================" << endl;

      // SIGNAL OK FOR ALL PDF
      RooRealVar mean3("mean3", "mean of gaussians", low, high);
      RooRealVar sigma3("sigma3", "width of gaussians", zpWidth*0.5, zpWidth*1.5);
      RooGaussian sgnl3("sgnl3", "Signal component", invMass, mean3, sigma3);
      mean3.setVal(masScan);
      mean3.setConstant(kTRUE);
      sigma3.setVal(zpWidth);
      //sigma3.setConstant(kTRUE);

      // HISTO PDF
      RooAddPdf model_3("model_3", "sigModel+bkgd shapes", RooArgList(sgnl3, fitExpoFun1), RooArgList(mu3));
      RooArgSet poi_3(mu3);
      ProfileLikelihoodCalculator plc_3(hist_data, model_3, poi_3);
      RooArgSet *nullParams_3 = (RooArgSet *)poi_3.snapshot();
      nullParams_3->setRealValue("mu3", 0);
      plc_3.SetNullParameters(*nullParams_3);
      HypoTestResult *htr_3 = plc_3.GetHypoTest();
      Double_t significance_histo = htr_3->Significance();
      if (isnan(significance_histo)) significance_histo = 0;
      if (isinf(significance_histo)) significance_histo = 0;
      cout << "fitRes:" << nCount << " " << masScan << " " << significance_histo << endl;
      yVal[nCount] = significance_histo;

      // plot signal candidates with background model and components for channel 2
      TCanvas *c3 = new TCanvas();
      RooPlot *xframe3 = invMass.frame();
      hist_data.plotOn(xframe3, DataError(RooAbsData::SumW2));
      model_3.plotOn(xframe3);
      model_3.plotOn(xframe3, Components(fitExpoFun1), LineStyle(kDashed), LineColor(kRed));
      hist_bkgd.plotOn(xframe3, DataError(RooAbsData::SumW2), MarkerColor(kGreen), LineColor(kGreen));
      xframe3->SetTitle("fit data + signal + background CH 1");
      xframe3->SetMinimum(1e-1);
      xframe3->SetMaximum(5e4);
      xframe3->Draw();
      c3->SetLogy();
      c3->Modified();
      c3->Update();
      c3->SaveAs(Form("zc_%d.pdf",nCount));
      delete c3;

   }

  TGraph *g_significance = new TGraph(nBin,xVal,yVal);
  TCanvas *c4 = new TCanvas();
  TH2F *h2 = new TH2F("h2","",400,low, high,100,0,20);
  h2->Draw();
  g_significance->SetLineColor(1);
  g_significance->SetMarkerColor(2);
  g_significance->SetMarkerStyle(20);
  g_significance->SetMarkerSize(1);
  g_significance->Draw("LP");

}