//Macro to use CTIDE EFractionFitter with an input of histograms

#ifndef __CINT__
#include "RooGlobalFunc.h"
#endif
#include "RooRealVar.h"
#include "RooDataSet.h"
#include "RooDataHist.h"
#include "RooGaussian.h"
#include "RooHistPdf.h"
#include "TCanvas.h"
#include "RooPlot.h"
#include "TTree.h"
#include "TH1D.h"
#include "TRandom.h"
#include <string>
#include <cmath>
#include <TH2.h>
#include <TStyle.h>
#include <iostream>

#include "EFractionFitter.h"
#include "EFractionFitter.C"
#include "EFitAnalysis.h" 

using namespace std;
using namespace RooFit ;


int use_RooFit() {

bool b_test = true;
bool b_FitOption = true; //Corresponds to boolean FractionFitter in EFitAnalysis above, false doesnt work...
bool b_makekmassplot = false; //make kaon mass ratio plot
bool b_makesampleplots = false; //make the sample real fake templates with prediction plots
	double ffake_sample = 0.2;

TFile file("hists_MC16_16Mar2019.root");
//TFile file("/eos/user/s/sferguso/KaonAnalysisData/MC16_test/MC16_hist.root"); //if using this change next line to .Get("hkmass");

TH1D *kmass = (TH1D*)file.Get("kmass");

TH1D *npix_MCreal = (TH1D*)file.Get("npix_MCreal");
TH1D *npix_MCfake = (TH1D*)file.Get("npix_MCfake");
TH1D *npix_side = (TH1D*)file.Get("npix_side");
TH1D *npix_sig = (TH1D*)file.Get("npix_sig");
TH1D *npix_sig_test = (TH1D*)file.Get("npix_sig_test");
TH1D *npix_side_test = (TH1D*)file.Get("npix_side_test");
//TH1D *npix_real_norm = (TH1D*)file.Get("npix_real_norm");

TH1D *nSCT_MCreal = (TH1D*)file.Get("nSCT_MCreal");
TH1D *nSCT_MCfake = (TH1D*)file.Get("nSCT_MCfake");
TH1D *nSCT_side = (TH1D*)file.Get("nSCT_side");
TH1D *nSCT_sig = (TH1D*)file.Get("nSCT_sig");
TH1D *nSCT_sig_test = (TH1D*)file.Get("nSCT_sig_test");
TH1D *nSCT_side_test = (TH1D*)file.Get("nSCT_side_test");
//TH1D *nSCT_real_norm = (TH1D*)file.Get("nSCT_real_norm");

TH1D *nTRT_MCreal = (TH1D*)file.Get("nTRT_MCreal");
TH1D *nTRT_MCfake = (TH1D*)file.Get("nTRT_MCfake");
TH1D *nTRT_side = (TH1D*)file.Get("nTRT_side");
TH1D *nTRT_sig = (TH1D*)file.Get("nTRT_sig");
TH1D *nTRT_sig_test = (TH1D*)file.Get("nTRT_sig_test");
TH1D *nTRT_side_test = (TH1D*)file.Get("nTRT_side_test");
//TH1D *nTRT_real_norm = (TH1D*)file.Get("nTRT_real_norm");

		int nbins_pix = npix_side->GetNbinsX();
		double_t pix_min;
		if (nbins_pix == 10) 	{
			 pix_min = -0.5;
		} else {
			 pix_min = 0.5;
		}
	TH1D *npix_real_norm = new TH1D("npix_real_norm","pixel hits data-driven real track distribution - normalized",nbins_pix,pix_min,9.5);
	npix_real_norm->Sumw2();
        TH1D *npix_fake_norm = new TH1D("npix_fake_norm","pixel hits data-driven fake track distribution - normalized",nbins_pix,pix_min,9.5);
        npix_fake_norm->Sumw2();
        TH1D *npix_sig_pred = new TH1D("npix_sig_pred","npix_sig_pred",10,-0.5,9.5);
        npix_sig_pred->Sumw2();

	TH1D *nSCT_real_norm = new TH1D("nSCT_real_norm","SCT hits data-driven real track distribution - normalized",18,-0.5,17.5);
	nSCT_real_norm->Sumw2();
        TH1D *nSCT_fake_norm = new TH1D("nSCT_fake_norm","SCT hits data-driven fake track distribution - normalized",18,-0.5,17.5);
        nSCT_fake_norm->Sumw2();
        TH1D *nSCT_sig_pred = new TH1D("nSCT_sig_pred","nSCT_sig_pred",18,-0.5,17.5);
        nSCT_sig_pred->Sumw2();

	TH1D *nTRT_real_norm = new TH1D("nTRT_real_norm","TRT hits data-driven real track distribution - normalized",70,-0.5,69.5);
	nTRT_real_norm->Sumw2();
        TH1D *nTRT_fake_norm = new TH1D("nTRT_fake_norm","TRT hits data-driven fake track distribution - normalized",70,-0.5,69.5);
        nTRT_fake_norm->Sumw2();
        TH1D *nTRT_sig_pred = new TH1D("nTRT_sig_pred","nTRT_sig_pred",70,-0.5,69.5);
        nTRT_sig_pred->Sumw2();

        double_t nevents_signal_region = npix_sig->Integral();
        double_t nevents_sideband_region = npix_side->Integral();
        double_t nevents_signal_region_test = npix_sig_test->Integral();
        double_t nevents_sideband_region_test = npix_side_test->Integral();

////Create Kaon Mass Plot and Fit

  //create roofit variable
        RooRealVar x("x","x",400,600) ;
  //create roofit gaussian
        RooRealVar mean("mean","mean",498,490,510) ;
        RooRealVar sigma("sigma","sigma",9,7,11) ;
        RooRealVar sig_yield("sig_yield","signal yield",10000,0,3000000);
        RooGaussian gauss("gauss","signal p.d.f.",x,mean,sigma) ;
  //create roofit quadratic
        RooRealVar c1("c1","coefficient of x term",-10, -400, -1) ;
        RooRealVar c2("c2","coefficient of x^2 term",10, 1, 30) ;
        RooRealVar c3("c3","coefficient of x^3 term",-0.01, -1, -0.0001) ;
        RooRealVar bkg_yield("bkg_yield","background yield",500000,0,4000000);
        RooPolynomial bkg("bkg","background", x, RooArgList(c1,c2,c3)) ;
  //create composite model model(x) = sig_yield*gauss(x) + bkg_yield*bkg(x)
        RooAddPdf model("model","model",RooArgList(gauss,bkg),RooArgList(sig_yield,bkg_yield)) ;
  //imposrt data from hist
        RooDataHist roodata("roodata","dataset with K0 mass",x,kmass) ;
  //fit model to data
        RooFitResult* r_model = model.fitTo(roodata,Extended(),Save());
  //roofit's version of a canvas:
        RooPlot* xframe = x.frame(Title("Kaon Mass with Gaussian Fit and Quadradic Background")) ;
        xframe->SetName("kaonmass_fit");
        xframe->SetXTitle("Kaon Mass (GeV)");
        xframe->SetYTitle("Number of Events");
  //plot
        roodata.plotOn(xframe) ;
        model.plotOn(xframe) ;
        model.plotOn(xframe,Components(bkg),LineColor(kGreen));


 //create integrals
        x.setRange("signal",mean.getVal()-2*sigma.getVal(),mean.getVal()+2*sigma.getVal());
        RooAbsReal* fsigregion_model = model.createIntegral(x,NormSet(x),Range("signal")); //this is the fraction of total events that are in the signal reg$
        RooAbsReal* fsigregion_bkg = bkg.createIntegral(x,NormSet(x),Range("signal")); //this is the fraction of background events that are in the signal re$
        //the number of event in the signal region that are actually signal (not background):
        Double_t nsigevents = fsigregion_model->getVal()*(sig_yield.getVal()+bkg_yield.getVal())-fsigregion_bkg->getVal()*bkg_yield.getVal();
        //calc error:
        Double_t error_nsigevents = sqrt(
                                                pow(fsigregion_model->getPropagatedError(*r_model)/fsigregion_model->getVal(),2)
                                                +(pow(sig_yield.getError(),2)+pow(bkg_yield.getError(),2))/pow(sig_yield.getVal()+bkg_yield.getVal(),2)
                                                + pow(fsigregion_bkg->getPropagatedError(*r_model)/fsigregion_bkg->getVal(),2)
                                                + pow(bkg_yield.getError()/bkg_yield.getVal(),2)
                                                )*nsigevents;

        //the fraction of events in the signal region that are actually signal (not background):
        Double_t fsig = nsigevents/(fsigregion_model->getVal()*(sig_yield.getVal()+bkg_yield.getVal()));
        //calc error:
        Double_t error_fsig = sqrt(
                                        pow(error_nsigevents/nsigevents,2)
                                        +pow(fsigregion_model->getPropagatedError(*r_model)/fsigregion_model->getVal(),2)
                                        +(pow(sig_yield.getError(),2)+pow(bkg_yield.getError(),2))/pow(sig_yield.getVal()+bkg_yield.getVal(),2)
                                        )*fsig;


cout << "fsigregion_model = " << fsigregion_model->getVal() << " +/- " << fsigregion_model->getPropagatedError(*r_model) << endl;
cout << "fsigregion_bkg = " << fsigregion_bkg->getVal() << " +/- " << fsigregion_bkg->getPropagatedError(*r_model) << endl;
cout << "nsigevents = " << nsigevents << " +/- " << error_nsigevents << endl;
cout << "fsig = " << fsig << " +/- " << error_fsig << endl;

//make kmass ratio plot
	if (b_makekmassplot) {
	TCanvas *cratio = new TCanvas("cratio","cratio",800,800);
        TPad *pad1 = new TPad("pad1", "pad1", 0, 0.3, 1, 1.0);
        pad1->SetBottomMargin(0); // Upper and lower plot are joined
        pad1->SetGridx();         // Vertical grid
        pad1->Draw();             // Draw the upper pad: pad1
        pad1->cd();               // pad1 becomes the current pad
        xframe->Draw();
        TLegend *l = new TLegend(0.1,0.6,0.4,0.9);
        std::string str2 = "fraction of in signal region = " +  std::to_string(fsig) + " +/- " + std::to_string(error_fsig);
        const char *chr2 = str2.c_str();
        l->AddEntry((TObject*)0,chr2,"");
        std::string str3 = "mean = " +  std::to_string(mean.getVal()) + " +/- " + std::to_string(mean.getError());
        const char *chr3 = str3.c_str();
        l->AddEntry((TObject*)0,chr3,"");
        std::string str4 = "sigma = " +  std::to_string(sigma.getVal()) + " +/- " + std::to_string(sigma.getError());
        const char *chr4 = str4.c_str();
        l->AddEntry((TObject*)0,chr4,"");
        std::string str5 = "num signal events = " +  std::to_string(sig_yield.getVal()) + " +/- " + std::to_string(sig_yield.getError());
        const char *chr5 = str5.c_str();
        l->AddEntry((TObject*)0,chr5,"");
        std::string str6 = "num background events = " +  std::to_string(bkg_yield.getVal()) + " +/- " + std::to_string(bkg_yield.getError());
        const char *chr6 = str6.c_str();
        l->AddEntry((TObject*)0,chr6,"");
        l->Draw();
        cratio->cd();          // Go back to the main canvas before defining pad2
        TPad *pad2 = new TPad("pad2", "pad2", 0, 0.05, 1, 0.3);
        pad2->SetTopMargin(0);
        pad2->SetBottomMargin(0.2);
        pad2->SetGridx(); // vertical grid
        pad2->Draw();
        pad2->cd();       // pad2 becomes the current pad
        TH1F *fitratio = new TH1F("fitratio","fitratio",100,400,600);
        RooDataHist* data_from_model = model.generateBinned(x,0,kTRUE);
        TH1F *kmass_fromroomodel = (TH1F*) data_from_model->createHistogram("kmass_fromroomodel",x,Binning(100,400,600));
        TH1F *kmass_fromroodata = (TH1F*) roodata.createHistogram("kmass_fromroodata",x,Binning(100,400,600));
        kmass_fromroomodel->Sumw2();
        kmass_fromroodata->Sumw2();
        fitratio->Sumw2();
        fitratio->Divide(kmass_fromroomodel,kmass_fromroodata,1,1);
        fitratio->GetYaxis()->SetTitle("ratio fit/data");
        fitratio->SetStats(kFALSE);
        fitratio->SetTitle("");
        fitratio->GetYaxis()->SetTitleSize(20);
        fitratio->GetYaxis()->SetTitleFont(43);
        fitratio->GetYaxis()->SetTitleOffset(1.55);
        fitratio->GetYaxis()->SetLabelFont(43); // Absolute font size in pixel (precision 3)
        fitratio->GetYaxis()->SetLabelSize(15);
        fitratio->GetXaxis()->SetTitleSize(20);
        fitratio->GetXaxis()->SetTitleFont(43);
        fitratio->GetXaxis()->SetTitleOffset(1.55);
        fitratio->GetXaxis()->SetLabelFont(43); // Absolute font size in pixel (precision 3)
        fitratio->GetXaxis()->SetLabelSize(15);
        fitratio->Draw();
        cratio->Print("PLOTS/kmass_ratio.pdf");
	}

// create the real templates: 

	npix_real_norm->Add(npix_sig,npix_side,1/nevents_signal_region,-1*(1-fsig)/nevents_sideband_region);
                for(int i=1;i<(nbins_pix+1);i++){
                if(npix_real_norm->GetBinContent(i)<0) npix_real_norm->SetBinContent(i,0);
                }
        double_t scale_npix_real_norm = npix_real_norm->Integral();
        npix_real_norm->Scale(1/scale_npix_real_norm);

        nSCT_real_norm->Sumw2();
        nSCT_real_norm->Add(nSCT_sig,nSCT_side,1/nevents_signal_region,-1*(1-fsig)/nevents_sideband_region);
                for(int i=1;i<19;i++){
                if(nSCT_real_norm->GetBinContent(i)<0) nSCT_real_norm->SetBinContent(i,0);
                }
        double_t scale_nSCT_real_norm = nSCT_real_norm->Integral();
        nSCT_real_norm->Scale(1/scale_nSCT_real_norm);

        nTRT_real_norm->Sumw2();
        nTRT_real_norm->Add(nTRT_sig,nTRT_side,1/nevents_signal_region,-1*(1-fsig)/nevents_sideband_region);
                for(int i=1;i<71;i++){
                if(nTRT_real_norm->GetBinContent(i)<0) nTRT_real_norm->SetBinContent(i,0);
                }
        double_t scale_nTRT_real_norm = nTRT_real_norm->Integral();
        nTRT_real_norm->Scale(1/scale_nTRT_real_norm);


// things that need to be declared before the loop:
	double_t f[1000] = { };
	double_t fpix[1000] = { };
	double_t rpix[1000] = { };
	double_t fpix_err[1000] = { };
        double_t rpix_err[1000] = { };
        double_t ffake;
	
	double_t chi2ndof_pix[1000] = { };

//start roofit for fake/real template fitting


// begin looping over fake fractions:
for(int a=1; a<1001; a++){

        ffake = .001*a; //initial guess at fake fraction
        f[a-1] = ffake;

// build arrays of fake templates and the corresponding predictions

	if(b_test || b_makesampleplots) ffake = ffake_sample;

        //fake templates:
        npix_fake_norm->Add(npix_side,npix_real_norm,1/nevents_sideband_region,-1*(1-ffake));
                for(int i=1;i<(nbins_pix+1);i++){
                if(npix_fake_norm->GetBinContent(i)<=0) npix_fake_norm->SetBinContent(i,0);
                }
        double_t scale_npix_fake_norm = npix_fake_norm->Integral(1,3);
        npix_fake_norm->Scale(1/scale_npix_fake_norm);

        nSCT_fake_norm->Add(nSCT_side,nSCT_real_norm,1/nevents_sideband_region,-1*(1-ffake));
                for(int i=1;i<19;i++){
                if(nSCT_fake_norm->GetBinContent(i)<=0) nSCT_fake_norm->SetBinContent(i,0);
                }
        double_t scale_nSCT_fake_norm = nSCT_fake_norm->Integral(3,12);
        nSCT_fake_norm->Scale(1/scale_nSCT_fake_norm);

        nTRT_fake_norm->Add(nTRT_side,nTRT_real_norm,1/nevents_sideband_region,-1*(1-ffake));
                for(int i=1;i<71;i++){
                if(nTRT_fake_norm->GetBinContent(i)<=0) nTRT_fake_norm->SetBinContent(i,0);
                }
        double_t scale_nTRT_fake_norm = nTRT_fake_norm->Integral(9,44);
        nTRT_fake_norm->Scale(1/scale_nTRT_fake_norm);
TH1F *h = new TH1F("h","h",nbins_pix,pix_min,9.5);
h->Add(npix_real_norm,npix_fake_norm,0.8,0.2);
TCanvas *c = new TCanvas("c","c",1);
h->Draw();
c->Print("PLOTS/h.pdf");

	RooRealVar npix("npix","npix",pix_min,9.5);
	RooDataHist npix_data("npix_data","npix_data",npix,h);
	RooRealVar f_real_pix("f_real_pix","f_real_pix",0.8,0,1);
	RooRealVar f_fake_pix("f_fake_pix","f_fake_pix",0.2,0,1);
	RooDataHist pix_real_temp("pix_real_data","pix_real_data",npix,npix_real_norm);
	RooHistPdf pix_Pdf_real("pix_Pdf_real","pix_Pdf_real",npix,pix_real_temp,0);

        RooDataHist pix_fake_temp("pix_real_data","pix_real_data",npix,npix_fake_norm);
	RooHistPdf pix_Pdf_fake("pix_Pdf_fake","pix_Pdf_fake",npix,pix_fake_temp,0);
	RooAddPdf pix_Pdf("pix_Pdf","pix_Pdf",RooArgList(pix_Pdf_real,pix_Pdf_fake),RooArgList(f_real_pix,f_fake_pix));
	RooFitResult *pix_result = pix_Pdf.fitTo(npix_data,Range(0.5,6.5));
	RooPlot *pixframe = npix.frame();

	npix_data.plotOn(pixframe);
	pix_Pdf.plotOn(pixframe);
	pix_Pdf.plotOn(pixframe,Components(pix_Pdf_fake),LineColor(kRed));

	rpix[a-1]=f_real_pix.getVal();
	fpix[a-1]=f_fake_pix.getVal();
        rpix_err[a-1]=f_real_pix.getError();
        fpix_err[a-1]=f_fake_pix.getError();

cout << "real fraction: " << f_real_pix.getVal() << " +/- " << f_real_pix.getError() << endl;
cout << "fake fraction: " << f_fake_pix.getVal() << " +/- " << f_fake_pix.getError() << endl;

	chi2ndof_pix[a-1]=pixframe->chiSquare();

	if(b_test){
	TCanvas *temp = new TCanvas("temp","temp",1);
	TPad *temp_pad1 = new TPad("temp_pad1", "temp_pad1", 0, 0.3, 1, 1.0);
        temp_pad1->SetBottomMargin(0); // Upper and lower plot are joined
        temp_pad1->SetGridx();         // Vertical grid
        temp_pad1->Draw();             // Draw the upper pad: temp_pad1
        temp_pad1->cd();  
	pixframe->Draw();
	temp->cd();
	TPad *temp_pad2 = new TPad("temp_pad2", "temp_pad2", 0, 0.05, 1, 0.3);
        temp_pad2->SetTopMargin(0);
        temp_pad2->SetBottomMargin(0.2);
        temp_pad2->SetGridx(); // vertical grid
        temp_pad2->Draw();
        temp_pad2->cd();       // temp_pad2 becomes the current pad
	TH1F *pixfitratio = new TH1F("pixfitratio","pixfitratio",nbins_pix,pix_min,9.5);
	RooDataHist *pix_data_from_Pdf = pix_Pdf.generateBinned(npix,0,kTRUE);
	TH1F *npix_fromPdf = (TH1F*) pix_data_from_Pdf->createHistogram("npix_fromPdf",npix,Binning(nbins_pix,pix_min,9.5));
	npix_fromPdf->Sumw2();
	pixfitratio->Sumw2();
	pixfitratio->Divide(npix_fromPdf,npix_side_test,1,1);
	pixfitratio->SetStats(kFALSE);
        pixfitratio->SetTitle("");
        pixfitratio->GetYaxis()->SetTitleSize(20);
        pixfitratio->GetYaxis()->SetTitleFont(43);
        pixfitratio->GetYaxis()->SetTitleOffset(1.55);
        pixfitratio->GetYaxis()->SetLabelFont(43); // Absolute font size in pixel (precision 3)
        pixfitratio->GetYaxis()->SetLabelSize(15);
        pixfitratio->GetXaxis()->SetTitleSize(20);
        pixfitratio->GetXaxis()->SetTitleFont(43);
        pixfitratio->GetXaxis()->SetTitleOffset(1.55);
        pixfitratio->GetXaxis()->SetLabelFont(43); // Absolute font size in pixel (precision 3)
        pixfitratio->GetXaxis()->SetLabelSize(15);
        pixfitratio->Draw();
	temp->Print("PLOTS/pixPdf_ratio.pdf");
	}

	if(b_test || b_makesampleplots) a=100000000;
    } //end of loop


return 1; 

}