#include <iostream>
#include <RooRealVar.h>
#include <RooWorkspace.h>
#include <RooExtendPdf.h>
#include <RooDataSet.h>
#include <RooStats/ModelConfig.h>
#include <RooStats/ProfileLikelihoodCalculator.h>
#include <RooStats/ProfileLikelihoodTestStat.h>
#include <RooStats/SamplingDistPlot.h>
#include <RooStats/LikelihoodInterval.h>
void excl_genD2nunu_2Dfit() {
 using namespace RooFit ;
 using namespace RooStats ;

 TChain chain1("Dsig");
 auto file = "/home/ckh/chkim/NtupleSET/sigNtuple2.root";
 chain1.Add(file);

 TChain chain3("Dsig");
 auto file_c = "/home/ckh/chkim/NtupleSET/ntuple2.root";
 chain3.Add(file_c);

 ROOT::EnableImplicitMT();
 ROOT::RDataFrame df1(chain1);
 ROOT::RDataFrame df3(chain3);

 TCut noTrk = "numTrk == 0";
 TCut fTagD = "sigDst_tagIsSignal != 1";
 TCut tTagD = "sigDst_tagIsSignal == 1";
 TCut noPiz = "numPI0 == 0";
 TCut noKLM = "numKL0 == 0";
 TCut noKS0 = "numKS0 == 0";
 TCut noLam0 = "numLam0 == 0";

 TCut wPiz = "numPI0 != 0";
 TCut wKLM = "numKL0 != 0";
 TCut wTrk = "numTrk != 0";
 TCut wKS0 = "numKS0 != 0";
 TCut wLam = "numLam0 != 0";

 TCut n2sTrk = "numTrk < 2";
 TCut n2sKL0 = "numKL0 < 2";
 TCut n2sKS0 = "numKS0 < 2";
 TCut n2sLm0 = "numLm0 < 2";
 TCut n2sPI0 = "numPI0 == 1";

 TCut fitRegion = "reco_M > 1.84 && reco_M < 1.875 && eECL < 2.1";
 TCut exclusiveD = noPiz && noKLM && noTrk && noKS0 && noLam0 && fitRegion;
 TCut SB = "(reco_M > 1.84 && reco_M < 1.859) || (reco_M > 1.871 && reco_M < 1.875)";
 TCut SR = "reco_M > 1.859 && reco_M < 1.871";
 TCut flatSel = SB && exclusiveD;
 TCut peakSel = SR && exclusiveD;

 auto df_ft = df3.Filter(flatSel.GetTitle());
 auto df_pk = df3.Filter(peakSel.GetTitle());

 TFile *f1 = new TFile("./PDFs/signal_pdf.root");
 RooWorkspace *w1 = (RooWorkspace *)f1->Get("w1");
 RooAbsPdf * sig = dynamic_cast<RooAbsPdf *>(w1->pdf("sig")->Clone());
 w1->loadSnapshot("mySnapshot");

 auto Df = df1.Filter(exclusiveD.GetTitle());

 auto gen_df = df3.Filter(exclusiveD.GetTitle());
 auto ex_h = Df.Histo1D({"h","exclusive D^{0} E_{ECL}; E_{ECL}(GeV); Events",42,0.0,2.1},"eECL");

 auto flat_h = df_ft.Histo1D({"h","exclusive D^{0} E_{ECL}; E_{ECL}(GeV); Events",42,0.0,2.1},"eECL");
 auto pk_h = df_pk.Histo1D({"h","exclusive D^{0} E_{ECL}; E_{ECL}(GeV); Events",42,0.0,2.1},"eECL");
 
 auto ex_hx = Df.Histo1D({"h","excluisive D^{0}; Mass(GeV); Events", 70, 1.84, 1.875},"reco_M");
 auto ex_gen_hx = gen_df.Histo1D({"h","excluisive D^{0}; Mass(GeV); Events", 70, 1.84, 1.875},"reco_M");
 auto ex_gen_hy = gen_df.Histo1D({"h","exclusive D^{0} E_{ECL}; E_{ECL}(GeV); Events",42,0.0,2.1},"eECL");

 // Declare observable x
 RooRealVar x("Mass","Mass(GeV)",1.84,1.875);
 RooRealVar y("eECL","E_{ECL} (GeV)", 0.0, 2.1);
 
 auto ex_rdhy = RooDataHist("ex_rdhy","ex_rdhy", RooArgList(y) ,ex_h.GetPtr(), 1.0);
 auto pk_rdhy = RooDataHist("pkF_rdhy","pk_rdhy", RooArgList(y) ,pk_h.GetPtr(), 1.0);
 auto ft_rdhy = RooDataHist("ftP_rdhy","ft_rdhy", RooArgList(y) ,flat_h.GetPtr(), 1.0);

 auto ex_rdhx = RooDataHist("ex_rdhx","ex_rdhx", RooArgList(x) ,ex_hx.GetPtr(), 1.0);

 auto gen_rdhx = RooDataHist("gen_rdhx", "gen_rdhx", RooArgList(x), ex_gen_hx.GetPtr(), 1.0);
 auto gen_rdhy = RooDataHist("gen_rdhy", "gen_rdhy", RooArgList(y), ex_gen_hy.GetPtr(), 1.0);


 RooRealVar fr1("fr1","fraction 1",0.4,0.0,1.0);
 RooRealVar m1("m1","mean of gaussians",1.8648,1.84,1.875);
 RooRealVar order("order","order of argus background", 0.6, 0.1, 3.5);
 
 RooRealVar ps1("ps1", "width of pbkg gaussians", 0.002, 0.0003, 0.015);
 RooRealVar ps2("ps2", "width of pbkg gaussians", 0.004, 0.0003, 0.015);
 RooRealVar ps3("ps3", "width of pbkg gaussians", 0.008, 0.0003, 0.015);

 RooGaussian pbkg1("pbkg1", "peaking background 1", x, m1, ps1);
 RooGaussian pbkg2("pbkg2", "peaking background 2", x, m1, ps2);
 RooGaussian pbkg3("pbkg3", "peaking background 3", x, m1, ps3);

 RooRealVar c1("c1","c1",3,1.875,4);
 RooRealVar c2("c2","c2",-0.4,-1.0,0.0);
 RooGenericPdf linear("linear","c1 + c2*(Mass)",RooArgSet(c1,c2,x));

 RooRealVar chiSq("chiSq","shape parameter",-150.0,-500.0,0.0);
 RooArgusBG argus("argus","argus PDF", x, RooConst(1.8712), chiSq, order);

 RooAddPdf flatC("flatC","flat Component",RooArgList(argus,linear),RooArgSet(fr1));

 // sum the signal components into a composite signal
 RooRealVar pkbg1frac("pkbg1frac","fraction of component 1 in peak bkg",0.4,0.,1.);
 RooRealVar pkbg2frac("pkbg2frac","fraction of component 1 in peak bkg",0.3,0.,1.);
 RooFormulaVar pkbg3frac("pkbg3frac","1 - pkbg1frac - pkbg2frac", "1 - @0 - @1",RooArgList(pkbg1frac,pkbg2frac));

 RooAddPdf pbkg("pbkg", "peak bkg", RooArgList(pbkg1,pbkg2,pbkg3), RooArgList(pkbg1frac,pkbg2frac,pkbg3frac));

 // Declare observable x

 // Creat Histogram PDF
 RooHistPdf ecl_eECL("ecl_eECL", "exclusive", y, ex_rdhy, 0);
 RooHistPdf ft_eECL("ft_eECL", "flat background", y, ft_rdhy, 0); // flat bkg from side band
 RooHistPdf sub_eECL("sub_eECL", "flat background", y, ft_rdhy, 0); // flat bkg from side band
 RooHistPdf sr_eECL("sr_eECL", "signal region background", y, pk_rdhy, 0); // peak + flat on signal region
 y.setBins(42);
 RooRealVar subR("subR","substraction ratio", -0.5, -0.75, -0.25);
 RooFormulaVar subL("subL","1 - subR","1 - @0",RooArgList(subR));

 RooAddPdf pk_eECL("pk_eECL","peak componenet", RooArgList(sub_eECL, sr_eECL), RooArgList(subR,subL));

 RooRealVar nsig("nsig", "number of signal events", -25, -100, 200);
 RooRealVar nbkg("nbkg", "number of background events", 4800, 1000, 6000);

 auto nEnt = chain3.GetEntries(exclusiveD);
 std::string nTevt = to_string(nEnt);
 std::string def_nbkg = nTevt + " - @0";

 RooRealVar bkFr("bkFr", "fraction of background", 0.1, 0.05, 0.8);

 RooProdPdf sig_2d("sig_2d", "exclusive D", *sig, ecl_eECL, 0.0);
 RooProdPdf pbkg_2d("pbkg_2d", "peak background", pbkg, pk_eECL, 0.0);
 RooProdPdf flat_2d("flat_2d", "flat background", flatC, ft_eECL, 0.0);

 RooAddPdf bkg_2d("bkg_2d","background",RooArgList(pbkg_2d,flat_2d),RooArgList(bkFr));
 RooAddPdf model_2d("model","sig + bkg",RooArgList(sig_2d,bkg_2d),RooArgList(nsig,nbkg));
 
 RooAddPdf bkgECL("bkgECL","bkg",RooArgList(pk_eECL,ft_eECL),RooArgList(bkFr));
 RooAddPdf ECL("ECL","eECL",RooArgList(ecl_eECL,bkgECL),RooArgList(nsig,nbkg));
 RooAddPdf bkgPDF("bkgPDF","bkg",RooArgList(pbkg,flatC),RooArgList(bkFr));
 RooAddPdf PDF("PDF","PDF",RooArgList(*sig,bkgPDF),RooArgList(nsig,nbkg));

 auto correlation_2D = new TH2F("correlation_2D", "eECL and recoil Mass of recoil D",70,1.84,1.875,42,0.0,2.1);
 chain3.Draw("eECL:reco_M >> correlation_2D",exclusiveD);
 TH2F * hh = (TH2F*) gDirectory->Get("correlation_2D");
 
 RooFitResult * fitR(nullptr);
 auto r2dh = RooDataHist("r2dh","r2dh", RooArgList(x,y), hh);
 
 // perform 2D fit and get result(fitR)
 fitR = model_2d.fitTo(r2dh,Extended(),RooFit::PrintLevel(-1),RooFit::Save());
 RooPlot *yframe = y.frame(Title("E_{ECL} (GeV)"), Bins(42));
 RooPlot *xframe = x.frame(Title("Mass (GeV)"), Bins(70));
 RooPlot *yframe2 = y.frame(Title("histogramPDF from signal MC sample"),Bins(42));

 ecl_eECL.plotOn(yframe2, LineColor(kBlue), LineStyle(kDotted));
 ft_eECL.plotOn(yframe2, LineColor(kRed), LineStyle(kDotted));
 pk_eECL.plotOn(yframe2, LineColor(kGreen), LineStyle(kDotted));

 gen_rdhx.plotOn(xframe);
 PDF.plotOn(xframe);
 model_2d.plotOn(xframe, Components(pbkg), LineColor(kGreen), LineStyle(kDotted));
 model_2d.plotOn(xframe, Components(flatC), LineColor(kRed), LineStyle(kDotted));
 model_2d.paramOn(xframe, Label(""), Format("NEU", AutoPrecision(1)), Layout(0.1,0.5,0.9), Parameters(RooArgSet(nsig,nbkg)));
 model_2d.Print("t");

 gen_rdhy.plotOn(yframe);
 ECL.plotOn(yframe);
 model_2d.plotOn(yframe, Components(pk_eECL), LineColor(kGreen), LineStyle(kSolid));
 model_2d.plotOn(yframe, Components(ft_eECL), LineColor(kRed), LineStyle(kSolid));
 model_2d.paramOn(yframe, Label(""), Format("NEU", AutoPrecision(1)), Layout(0.9,0.6,0.9), Parameters(RooArgSet(nsig,nbkg)));
 model_2d.Print("t");

 auto Cx = new TCanvas();
 auto Cy = new TCanvas();
 auto nC = new TCanvas();

 nC -> cd();
 yframe2 -> Draw();
 nC -> SaveAs("eECL-histogramPDF-D2nunu.png");
 nC -> Update();
 
 Cx -> cd();
 xframe -> Draw();
 Cx -> SaveAs("exclusive-D2nunu-genBKG_2dFit_Mass.png");
 Cx -> Update();

 Cy -> cd();
 yframe -> Draw();
 Cy -> SaveAs("exclusive-D2nunu-genBKG_2dFit_eECL.png");
 Cy -> Update();
 std::cout << gen_df.Count().GetValue() << std::endl;

 TString filename = "./PDFs/UL_test.root";
 TFile fileA(filename);
 if (!fileA.IsOpen()) {
   std::cerr << "Error: Could not open file " << filename << std::endl;
   return;
 }

 // TTree
 TTree* tree = dynamic_cast<TTree*>(fileA.Get("Dsig"));
 if (!tree) {
   std::cerr << "Error: Could not find tree Dsig in file " << filename << std::endl;
   return;
 }
 
 // RooRealVar
 RooRealVar reco_M("reco_M", "Mass(GeV)", 1.84, 1.875);
 RooRealVar eECL("eECL", "E_{ECL} (GeV)", 0.0, 2.1);

 RooDataSet * data = new RooDataSet("data", "dataset", tree, RooArgSet(reco_M, eECL));
 fileA.Close();

 RooWorkspace w6("w6");

 m1.setConstant(kTRUE); ps1.setConstant(kTRUE); ps2.setConstant(kTRUE); ps3.setConstant(kTRUE);
 c1.setConstant(kTRUE); c2.setConstant(kTRUE); chiSq.setConstant(kTRUE); order.setConstant(kTRUE);
 pkbg1frac.setConstant(kTRUE); pkbg2frac.setConstant(kTRUE); 
 bkFr.setConstant(kTRUE); 
 fr1.setConstant(kTRUE);
 subR.setConstant(kTRUE);
 nbkg.setConstant(kTRUE);
 
 w6.import(model_2d);
 w6.import(*data);

 ModelConfig mc("modelConfig", &w6);
 mc.SetPdf(model_2d);
 mc.SetParametersOfInterest(RooArgSet(nsig));
 mc.SetNuisanceParameters(RooArgSet(m1,ps1,ps2,ps3,c1,c2,chiSq,order,pkbg1frac,pkbg2frac,fr1,bkFr,subR,nbkg)); 
 mc.SetObservables(RooArgSet(reco_M,eECL));
 w6.import(mc);

 // get nll function from fitted 2D PDF function
 auto nll = model_2d.createNLL(*data);
 auto nsig_nll = nll->createProfile(nsig);

 RooPlot *nll_frame = nsig.frame(Title("nsig"),Bins(100),Range(-100,200));
 auto Cnll = new TCanvas("nll","Negative Log-Likelihood", 900, 800);
 nsig_nll->plotOn(nll_frame, ShiftToZero(), LineColor(kRed));

 Cnll->cd();
 nll_frame->Draw(); 
 Cnll->SaveAs("nsig_NLL.png");
 Cnll->Close();
 
 // estimation nsig Upper limit based using NLL
 ProfileLikelihoodCalculator plc(*data, mc);
 plc.SetParameters(nsig);
 plc.SetConfidenceLevel(0.90);

 LikelihoodInterval *interval = plc.GetInterval();
 auto UL_val = interval->UpperLimit(nsig);

 std::cout << "fitResult : " << std::endl;
 fitR->Print("v");
 std::cout << "Upper Limit on Signal Yield at 90% CL: " << UL_val << std::endl;
}