// implementing s_weight to plot comparison plots for the control channel.
//
// \author Rishabh Raturi

#include "RooRealVar.h"
#include "RooStats/SPlot.h"
#include "RooDataSet.h"
#include "RooRealVar.h"
#include "RooGaussian.h"
#include "RooExponential.h"
#include "RooChebychev.h"
#include "RooAddPdf.h"
#include "RooProdPdf.h"
#include "RooAddition.h"
#include "RooProduct.h"
#include "TCanvas.h"
#include "RooAbsPdf.h"
#include "RooFit.h"
#include "RooFitResult.h"
#include "RooWorkspace.h"
#include "RooConstVar.h"
#include <iomanip>

using namespace RooFit;
using namespace RooStats;

void AddModel(RooWorkspace *);
void AddData(RooWorkspace *);
void DoSPlot(RooWorkspace *);

void splot()
{
	//create a workspace
	RooWorkspace *wspace = new RooWorkspace("w_space");

	//create a function to add the real-time data into the workspace
        AddData(wspace);

	//add signal and background model to the workspace
	// create a function AddModel() which has the description of the model
	AddModel(wspace);

	//choice to print the workspac
	// wspace->Print();
	
	//this function creates a new dataset with s_weight applied to every event
	DoSPlot(wspace);

	//this fucntion plots the comparison plots for the control channel
	//MakePlots(wspace);

	//save workspace in a file
	wspace->writeToFile("workspace.root");
	
}

void AddModel(RooWorkspace *ws)
{
	//make model for JPsi - double crystal ball for signal and exponential for the background
	TChain *trd = new TChain("tree");
        trd->Add("/home/rishabh/project/root_files/val/2016/red_ntuples/JPsi_Data_ForFit.root");


	double bm_min(5.2), bm_max(5.6);
	RooRealVar Bmass("Bmass","#bf{m(K^{+}K^{-}#mu^{+}#mu^{-}) [GeV]}", bm_min, bm_max);	
	cout << "make signal model" << endl;
	RooRealVar  mean("mean","common means for Crystal Balls", 5.367, 5.2, 5.6);
        RooRealVar  sigma1("sigma1","sigma of CB1",  0.02445);
        RooRealVar  sigma2("sigma2","sigma of CB2",  0.04508);
        RooRealVar  sigM_frac("sigM_frac","fraction of CB",  0.51306);
        RooRealVar  n1("n1", "", 1.9277906);
        RooRealVar  n2("n2", "", 3.6134270);
        RooRealVar  alpha1("alpha1","alpha for CB1", 1.9130096);
        RooRealVar  alpha2("alpha2","alpha for CB2", -1.9411376);

	RooCBShape CB1("CB1","Crystal Ball-1", Bmass, mean, sigma1, alpha1,n1);
        RooCBShape CB2("CB2","Crystal Ball-2", Bmass, mean, sigma2, alpha2,n2);

	RooAddPdf CB("CB","CB1+CB2", RooArgList(CB1,CB2), RooArgList(sigM_frac));

	RooRealVar lambda("lambda","slope",-10.,0.);
        RooExponential bkgE("bkgE","exponential PDF", Bmass, lambda);

	RooRealVar nsig("nsig","nsig",1E3,0,1E9);
        RooRealVar nbkgE("nbkgE","number of exponential bkg events", 8000, 0, 1E7);

	//RooRealVar nbkgE("nbkgE","number of exponential bkg events", 8000, 0, 1E7);

	//defining pdf for the peaking background- two crystal balls
        RooRealVar  m("mn","common means for Crystal Balls", 5.406, 5.3, bm_max);
        RooRealVar  s1("s1","sigma of CB1", 0.0663832 );
        RooRealVar  s2("s2","sigma of CB2", 0.0437879 );
        RooRealVar  frac("sigM_frac","fraction of CB", 0.4162298);
        RooRealVar  n_1("n1", "", 25.9750365);
        RooRealVar  n_2("n2", "", 99.8215927);
        RooRealVar  a1("a1","alpha for CB1",  1.1403154 );
        RooRealVar  a2("a2","alpha for CB2", -0.4118587 );

        RooCBShape PB1("PB1", "", Bmass, m, s1, a1, n_1);
        RooCBShape PB2("PB2", "", Bmass, m, s2, a2, n_2);

	//combining crystal balls for peaking background
        RooAddPdf PB("PB", "PB1+PB2", RooArgList(PB1, PB2), RooArgList(frac));
	
	RooRealVar nsigD("nsigD","nsigD",1E3,500,1.1*trd->GetEntries());
	RooRealVar fy("fy", "fy", 0.06903928336);
	RooProduct nsigpb("nsigpb", "nsigpb", RooArgList(nsigD, fy));

	//create final model for fitting
        RooAddPdf model("model", "CB+PB+e", RooArgList(CB, PB, bkgE), RooArgList(nsig, nsigpb, nbkgE));
        model.Print("t");

	std::cout << "import model to workspace" << std::endl;
	ws->import(model);
}

void AddData(RooWorkspace *ws)
{
	// this function add the real-time data for phimumu
	
	TChain *ch = new TChain("tree");
        ch->Add("/home/rishabh/project/root_files/val/2016/red_ntuples/JPsi_Data_ForFit.root");
        TTree *tr = ch;

        int nentries_ = tr->GetEntries();
        cout << "\n=> total entries in data tree = " << nentries_ << endl;

	double bm_min(5.2), bm_max(5.6);

	RooRealVar Bmass("Bmass","#bf{m(K^{+}K^{-}#mu^{+}#mu^{-}) [GeV]}", bm_min, bm_max);
        /*RooRealVar Mumumass("Mumumass","M^{#mu#mu}", 0., 10.);
        RooRealVar Mumumasserr("Mumumasserr","Error of M^{#mu#mu}", 0., 10.);
        RooRealVar Phimass("Phimass","phi mass", 0.,10.);
        RooRealVar Kmpt("Kmpt","K- pt", 0., 200.);
        RooRealVar Kppt("Kppt","K+ pt", 0., 200.);
        RooRealVar Kmtrkdcasigbs("Kmtrkdcasigbs","K^{-} track DCA/#sigma beam spot", 0., 1000.);
        RooRealVar Kptrkdcasigbs("Kptrkdcasigbs","K^{+} track DCA/#sigma beam spot", 0., 1000.);
        RooRealVar Blxysig("Blxysig","Blxy sig.", 0., 1000.);
        RooRealVar Bcosalphabs2d("Bcosalphabs2d","Bcosalphabs2d", 0., 1.);
        RooRealVar Bvtxcl("Bvtxcl","B vtx cl.", 0., 1.);*/

	RooArgSet  observables(Bmass);

	RooDataSet data("data","dataset with Bmass", ch, observables);

	TCut c1 = Form("Bmass>%f && Bmass<%f",bm_min,bm_max);
        TCut cutTotal = c1  ;

	RooDataSet *redData = (RooDataSet*)data.reduce(cutTotal);   //redData consists data to be fitted

	//importing this dataset to the workspace
	ws->import(*redData, Rename("data"));

	
}

void DoSPlot(RooWorkspace *ws)
{
	TChain *ch = new TChain("tree");
   	ch->Add("/home/rishabh/project/root_files/val/2016/red_ntuples/JPsi_Data_ForFit.root");
	TTree *tr = ch;
	std::cout << "calculate sWeight" << std::endl;

	// get what you need out of the workspace
	RooAbsPdf *model = ws->pdf("model");
	RooRealVar *nsig =  ws->var("nsig");
	RooRealVar *nbkgE = ws->var("nbkgE");
	RooDataSet *data = (RooDataSet *)ws->data("data");

	//fit model to data
	model->fitTo(*data);

	RooMsgService::instance().setSilentMode(true);

	// now make use of the SPlot class to add s_weight to out dataset
	// based on out model and yields
	RooStats::SPlot *sData = new RooStats::SPlot("sData", "An SPlot", *data, model, RooArgList(*nsig, *nbkgE));

	// create new root file for s_weighted output
    TFile *outFile = new TFile("JPsi_DataSet_with_sWeight.root", "RECREATE");
    outFile->cd(); 

    //cloning original tree
    TTree *tr1 = tr->CloneTree(); 

    double s_wt;
    double b_wt;

    TBranch *_s_wt = tr1->Branch("s_wt", &s_wt, "s_wt/D");
    TBranch *_b_wt = tr1->Branch("b_wt", &b_wt, "b_wt/D");

	// check that out weigt have the desired property
	std::cout << "Check SWeights:" << std::endl;

	std::cout << std::endl
		  << "N_{sig} " << nsig->getVal() << ": From s_weight N_{sig}"
		  << sData->GetYieldFromSWeight("nsig") <<std::endl;

	std::cout << std::endl
                  << "N_{bkg} " << nbkgE->getVal() << ": From s_weight N_{bkg}"
                  << sData->GetYieldFromSWeight("nbkgE") <<std::endl;

	for (Int_t i = 0; i < tr1->GetEntries(); i++)
	 {
		if (i%1000 == 0) std::cout << "--- ... Processing event: " << i << std::endl;
      	s_wt = sData->GetSWeight( i, "nsig");
      	b_wt = sData->GetSWeight( i, "nbkgE");

      	_s_wt->Fill();
     	_b_wt->Fill();
    
     }

	 std::cout << "Loop Finished" << std::endl;


	tr1->Write();
    outFile->Write();
    std::cout << "Files written" << std::endl;
    outFile->Close();
}