// This code performs extended likelihood fit on data sample, fixing signal shapes from simulation
// samples. Option to toggle between ranged fit and full fit.

#include <iostream>
using namespace std;
using namespace RooFit;

RooDataSet *GetNewWeightedMCDataset(TTree *&tree, RooArgSet &mc_arg_set, TString &CutList, TString lname);
void FitSignalPdf(RooAddPdf*& pdf, RooDataSet *&mcdata, RooCmdArg Range, TString lname = "varName");
void SetConstant(RooArgSet *& args, Bool_t value=true);
RooWorkspace* GetRooWorkspace();

RooRealVar *Bmass      = new RooRealVar("Bmass","B^{0} mass", 5.1, 5.6,"GeV");
RooRealVar *NewWeight  = new RooRealVar("NewWeight", "NewWeight", 0, 20);
TString BDT     	   = "Bdt2_";    // Name of BDT variable
TString ReweightName   = "Reweight"; // 2nd weight variable 

// Toggle between ranged-fit and full fit. 
bool rangeFlag = true; // make it false to fit in full range

void RangedFit(){
	RooMsgService::instance().setGlobalKillBelow(RooFit::ERROR);
    Bmass->setRange("sigRegion1", 5.1, 5.33);
    Bmass->setRange("sigRegion2", 5.41, 5.6);  
	TString outDir 	  = ".";
    TString path = "/afs/cern.ch/user/p/pkalbhor/public/final_weighted_tobecomp/";
	TString mcinput   = path + "sel_2018BdMC_wiso_Minibdt_Presel_mc_cut_bdt_s0.root";
	TString BsMCinput = path + "sel_2018BsMC_wiso_Minibdt_Presel_mc_cut_bdt_s0.root";
	TString datainput = path + "sel_Charmonium_2018data_wiso_Minibdt_Presel_data_cut_bdt_s0.root";

	TString AllSelection = "("+BDT+" > 0.57)";

    // INPUT
    TFile *tempfile = new TFile("temp.root","RECREATE");
    TFile *dataFile = new TFile(datainput);
    TTree *dataTree = (TTree*)dataFile->Get("tree");
    cout << "Input Data File: " << datainput << endl;

    TFile *BsMCFile = new TFile(BsMCinput);
    TTree *BsMCTree = (TTree*)BsMCFile->Get("tree");    
    cout << "Input Bs MC File: " << BsMCinput << endl;

    TFile *mcFile   = new TFile(mcinput);
    tempfile->cd();
    TTree *BdMCTree   = (TTree*)mcFile->Get("tree");
    cout << "Input Bd MC File: " << mcinput << endl;    

    // Variables for RooArgSet
    RooRealVar *Biso  	   = new RooRealVar("Biso57","Biso57", 0.2, 1.05, "");
    RooRealVar *Bdt 	   = new RooRealVar(BDT,"Bdt", 0.05, .9, "");
    RooRealVar *fpuw8      = new RooRealVar("fpu_w8","fpuw8", 1, 0, 10, "");
    RooRealVar *Reweight   = new RooRealVar(ReweightName,"Reweight", 1, 0, 10, "");

    RooArgSet mc_arg_set  = RooArgSet(*Bmass, *Biso, *Bdt, *fpuw8, *Reweight);
    RooArgSet data_arg_set= RooArgSet(*Bmass, *Biso, *Bdt);
    RooDataSet *Bd_mcData = GetNewWeightedMCDataset(BdMCTree, mc_arg_set, AllSelection, "Bd_forBiso");
    RooDataSet *Bs_mcData = GetNewWeightedMCDataset(BsMCTree, mc_arg_set, AllSelection, "Bs_forBiso");
    RooDataSet *data      = new RooDataSet("Dataset", "Dataset for data", data_arg_set, Import(*dataTree), Cut(AllSelection));

    // Get PDF
    RooWorkspace *wspace = GetRooWorkspace();
    RooAddPdf *massPdf = (RooAddPdf*)wspace->obj("mass_pdf");
    RooAddPdf *BdPdf   = (RooAddPdf*)wspace->obj("BdPdf");
    RooAddPdf *BsPdf   = (RooAddPdf*)wspace->obj("BsPdf");

    RooCmdArg range    = (rangeFlag) ? Range("sigRegion1,sigRegion2") : RooCmdArg::none();
    RooCmdArg cutRange = (rangeFlag) ? CutRange("sigRegion1,sigRegion2") : RooCmdArg::none();

    // Fit Bd Signal Mass and fixed its parameters
    FitSignalPdf(BdPdf, Bd_mcData, range);
	// Fit Bs Signal Mass and fixed its parameters
    FitSignalPdf(BsPdf, Bs_mcData, range);


    // Perform a fit and plot
    cout<< "---------------------------------Performing data fit-------------------------------------" << endl; 
    RooFitResult *fitResult = massPdf->fitTo(*data, Extended(), Strategy(2), Save(1), PrintLevel(0), range);
    fitResult->Print("v");

    // Plot PDF and data 
    TCanvas *c2  	= new TCanvas("c2", "c2");
	gPad->SetLogy(1);
    RooPlot *frame  = Bmass->frame();
    data->plotOn(frame);
    massPdf->plotOn(frame);
    massPdf->plotOn(frame, Components("BdPdf"), LineStyle(kDashed), LineColor(kGreen));
    massPdf->plotOn(frame, Components("BsPdf"), LineStyle(kDashed), LineColor(kOrange));
    massPdf->plotOn(frame, Components("mass_bkg"), LineStyle(kDashed), LineColor(kRed));
    massPdf->plotOn(frame, Components("pdf_errf"), LineStyle(kDashed), LineColor(kMagenta));
    frame->Draw();
    c2->SaveAs("Bmass.pdf");
}

// Adding 2 weights to RooDataSet 
RooDataSet *GetNewWeightedMCDataset(TTree *&tree, RooArgSet &mc_arg_set, TString &CutList, TString lname){
    cout<< ">> Obtaining RooDataSet with NewWeight as weight variable." << endl;

    RooDataSet *tempData = new RooDataSet("tempDataset"+lname, "tempDataset", mc_arg_set, Import(*tree), Cut(CutList));
    
    RooArgSet CopiedArgSet = (RooArgSet)mc_arg_set;
    CopiedArgSet.add(*NewWeight);
    RooDataSet *finalDataset = new RooDataSet("finalDataset"+lname,"finalDatset"+lname, CopiedArgSet, WeightVar("NewWeight"));
    for(Int_t j=0; j<tempData->numEntries(); j++){
        RooArgSet *args = (RooArgSet*)tempData->get(j);
        Float_t myweight = ((RooRealVar*)args->find("fpu_w8"))->getVal()*((RooRealVar*)args->find(ReweightName))->getVal();
        TIterator *args_it = args->createIterator();
        RooRealVar *arg= (RooRealVar*)args_it->Next();
        TIterator *column_args_it = CopiedArgSet.createIterator();
        RooRealVar *column_arg = (RooRealVar*)column_args_it->Next();
        while(arg){
            column_arg->setVal(arg->getVal());
            arg = (RooRealVar*)args_it->Next();
            column_arg = (RooRealVar*)column_args_it->Next();
        }
        ((RooRealVar*)CopiedArgSet.find("NewWeight"))->setVal(myweight);
        finalDataset->add(CopiedArgSet, myweight);
    }
    return finalDataset;
}

void FitSignalPdf(RooAddPdf*& pdf, RooDataSet *&mcdata, RooCmdArg Range, TString lname = "varName"){
    cout<< "---------------------------------Performing "<< pdf->GetTitle() <<" fit-------------------------------------" << endl; 
    RooArgSet *args = pdf->getParameters(mcdata);
    SetConstant(args, false);
    RooFitResult *fitResult = pdf->fitTo(*mcdata, SumW2Error(true), Strategy(2), Save(), PrintLevel(-3), Range);
    cout<<"Signal Fit: "<<endl; fitResult->Print("v");
    SetConstant(args, true);
}

// Set all variables in RooArgSet to Constant
void SetConstant(RooArgSet *& args, Bool_t value=true){
    TIterator *args_it = args->createIterator();
    RooRealVar *arg= (RooRealVar*)args_it->Next();
    while(arg){
        arg->setConstant(value);
        arg = (RooRealVar*)args_it->Next();
    }
}

RooWorkspace* GetRooWorkspace(){
	RooMsgService::instance().setGlobalKillBelow(RooFit::ERROR);

	RooRealVar frac1("frac1","Fraction1 for Bd signal", 0.658697, 0., 1.);
	RooRealVar frac2("frac2","Fraction2 for Bd signal", 3.20718e-01, 0., 1.);
	RooRealVar Bdmean("Bdmean","Bd mass mean", 5.27978, 5.25, 5.30);
	RooRealVar sigma1("sigma1","Sigma1 for Bd", 9.10867e-03, 0.0001, 5);
	RooRealVar sigma2("sigma2","Sigma2 for Bd", 2.47588e-02, 0.0001, 5);
	RooRealVar sigma3("sigma3","Sigma3 for Bd", 8.93581e-02, 0.0001, 5);
	RooGaussian B_G1("B_G1", "B0 Gauss1", *Bmass, Bdmean, sigma1);
	RooGaussian B_G2("B_G2", "B0 Gauss2", *Bmass, Bdmean, sigma2);
	RooGaussian B_G3("B_G3", "B0 Gauss2", *Bmass, Bdmean, sigma3);
	RooAddPdf BdPdf("BdPdf", "Bd Mass Signal", RooArgList(B_G1, B_G2, B_G3), RooArgList(frac1, frac2));

	RooRealVar Bsmean("Bsmean","Bs mass mean", 5.367, 5.33, 5.39);
	RooRealVar s1("s1","Sigma1 for Bs", 0.01, 0.0001, 5);
	RooRealVar s2("s2","Sigma2 for Bs", 0.02, 0.0001, 5);
    RooRealVar s3("s3","Sigma3 for Bs", 0.03, 0.0001, 5);
	RooRealVar frac3("frac3","Fraction3 for Bs", 0.65, 0., 1.);
    RooRealVar frac4("frac4","Fraction4 for Bs", 0.3, 0., 1.);
	RooGaussian G1("G1", "Gauss1", *Bmass, Bsmean, s1);
	RooGaussian G2("G2", "Gauss2", *Bmass, Bsmean, s2);
    RooGaussian G3("G3", "Gauss3", *Bmass, Bsmean, s3);    
    RooAddPdf BsPdf("BsPdf", "Bs mass pdf", RooArgList(G1, G2, G3), RooArgList(frac3, frac4));	

	RooRealVar bkgSlope("mass_bkgSlope", "bkg slope", -2.3626e+00, -100, 100);//
	RooExponential bkgPdf("mass_bkg", "mass bkg", *Bmass, bkgSlope);

	RooRealVar div("div","", 1.9042e-02, 0.0001, 1.8);
	RooRealVar shoulder("shoulder","", 5.1392e+00, 5.1, 5.22);
	RooGenericPdf pdf_errf("pdf_errf","", "( TMath::Erf( (-@0 + @1)/@2 ) + 1 )",RooArgSet(*Bmass,shoulder,div));

	RooRealVar nSig("nSig", "Number of Signal Events", 3.0664e+04, 0., 1e+07);
	RooRealVar nBkg("nBkg", "Number of Comb Bkg events", 2.7834e+03, 0., 1e+07);
	RooRealVar nPeak("nPeak", "Number of Bs Bkg events", 2.9491e+02, 0., 500);
	RooRealVar nErrf("nErrf", "Number of ErrF Bkg events", 1.7073e+03, 0., 1e+04);

    auto *massPdf = new RooAddPdf("mass_pdf", "Total Bmass PDF", RooArgList(BdPdf, BsPdf, bkgPdf, pdf_errf), RooArgList(nSig, nPeak, nBkg, nErrf));
    auto *wspace = new RooWorkspace("wspace");
    wspace->import(*massPdf);
    return wspace;
}