#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 UL_D2nunu() {
 using namespace RooFit;
 using namespace RooStats;
 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 fitRegion = "reco_M > 1.84 && reco_M < 1.875 && eECL < 2.1";
 TCut PeakBkg = fTagD && fitRegion && (wPiz || wKLM || wTrk || wKS0 || wLam);
 TCut FlatBkg = tTagD && fitRegion && (wPiz || wKLM || wTrk || wKS0 || wLam);
 TCut exclusiveD = noPiz && noKLM && noTrk && noKS0 && noLam0 && fitRegion;

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

 // TTree
 TTree* tree = dynamic_cast<TTree*>(file.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
 RooDataSet * data = new RooDataSet("data", "dataset", tree, RooArgSet(reco_M, eECL));

 RooWorkspace w("w");
 TFile *f5 = new TFile("./PDFs/D2nunu_sig_2Dpdf.root");
 RooWorkspace *w5 = (RooWorkspace *)f5->Get("w5");
 RooAbsPdf *sig_2d = dynamic_cast<RooAbsPdf*>(w5->pdf("sig_2d")->Clone());
 w5->loadSnapshot("mySnapshot");
 TFile *f4 = new TFile("./PDFs/D2nunu_bkg_2Dpdf.root");
 RooWorkspace *w4 = (RooWorkspace *)f4->Get("w4");
 RooAbsPdf *bkg_2d = dynamic_cast<RooAbsPdf*>(w4->pdf("bkg_2d")->Clone());
 w4->loadSnapshot("mySnapshot");

 w.import(*sig_2d);
 w.import(*bkg_2d);

 RooRealVar nsig("nsig", "signal yield", 0, 0, 200);
 RooRealVar nbkg("nbkg", "background yield", 3000, 0, 4500);

 RooExtendPdf esig("esig", "extended signal pdf", *sig_2d, nsig);
 RooExtendPdf ebkg("ebkg", "extended background pdf", *bkg_2d, nbkg);

 RooAddPdf model("model", "signal + background", RooArgList(esig,ebkg));

 ModelConfig mc("ModelConfig", &w);
 mc.SetPdf(model);
 mc.SetParametersOfInterest(RooArgSet(nsig));
 mc.SetObservables(RooArgSet(reco_M, eECL));
 mc.SetSnapshot(RooArgSet(nsig));

 ProfileLikelihoodCalculator plc(*data, mc);
 plc.SetParameters(nsig);
 plc.SetConfidenceLevel(0.90);

 LikelihoodInterval *interval = plc.GetInterval();
 std::cout << "Upper Limit on Signal Yield at 90% CL: " << interval->UpperLimit(nsig) << std::endl;

}