#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>
#include <iostream>
#include <string>
#include <sstream>

using namespace RooFit;
using namespace RooStats;

double bm_min(5.27), bm_max(5.6);
double ph_min(1.525-0.079), ph_max(1.525+0.079);
RooRealVar Bmass("Bmass","#bf{m(K^{+}K^{-}#mu^{+}#mu^{-}) [GeV]}", bm_min, bm_max);
RooRealVar Phimass("Phimass", "#bf{m(K^{+}K^{-}) GeV}", ph_min, ph_max);

void sWeightsBsToJpsif2p(const char* mode, const char* year){

    gROOT->SetBatch(1);
    const char* file = Form("../../red_ntuples/%s", year);
    const char* filename = Form("red_%s_Data_cutopt_bdtr.root", mode);

    cout << file << "/" << filename << endl;

    TChain *ch = new TChain("tree");
    ch->Add(Form("%s/%s", file, filename));
    TTree* tr = ch;

    cout << "\n Entries: " << tr->GetEntries() << endl;

    TFile* fin = new TFile(Form("../fits/wspaces/ws_%s_2DfitResults_%s.root", mode, year));
    RooWorkspace* w3 = (RooWorkspace*)fin->Get("w2");

    // fit results
    RooFitResult* fitres=(RooFitResult*)w3->genobj(Form("fitresult_model_%s_%s_data", mode, year));
    // Dataset
    RooDataSet *redData =(RooDataSet*)w3->data("data");
    // Yield paramters
    RooAbsArg* nsigDarg = fitres->floatParsFinal().find(Form("nsigD_%s_%s", mode, year));
    RooRealVar* nsigD = dynamic_cast<RooRealVar*>(nsigDarg);
    RooAbsArg* nbkgEarg = fitres->floatParsFinal().find(Form("nbkgE_%s_%s", mode, year));
    RooRealVar* nbkgE = dynamic_cast<RooRealVar*>(nbkgEarg);
    // fit model
    RooAbsPdf* model = w3->pdf(Form("model_%s_%s", mode, year));

    if (nsigD) nsigD->setRange(0, nsigD->getVal() * 2);  // 
    if (nbkgE) nbkgE->setRange(0, nbkgE->getVal() * 2);

    fitres->Print("v");

    fitres->correlationMatrix().Print();
    /*fitres->covarianceMatrix().Print();
    return;*/

    TCanvas *c2 = new TCanvas("c","c",800, 700);
    RooPlot *xframe = Bmass.frame(Title(""), Bins(160));
     xframe->SetTitle("");
    xframe->GetXaxis()->SetTitle("#bf{m(K^{+}K^{-}#mu^{+}#mu^{-}) [GeV]}");

    redData->plotOn(xframe,RooFit::Name("data"));
    model->plotOn(xframe, 
                 RooFit::Name("Full PDF"), 
                 LineColor(kBlue));
    model->plotOn(xframe, 
                 RooFit::Name("CB"),         
                 Components(Form("CB_%s_%s", mode, year)), 
                 LineStyle(kDashed), 
                 LineColor(kGreen));

    model->plotOn(xframe,
                  RooFit::Name("bkgE_data"),
                  Components(Form("bkgE_%s_%s", mode, year)),
                  LineColor(kRed),
                  DrawOption("F"),
                  FillColor(kRed),
                  FillStyle(3003));

    cout << "Signal Yield: " << nsigD->getVal() << " +/- " << nsigD->getError() << endl;
    std::cout << "Integral of model: " << model->getNorm() << std::endl;

    xframe->Draw();
    c2->SaveAs(Form("../plots/%s/sWeight/%s/BmassDataFit_sWeight_BsToJpsif2p_%s.pdf", year, mode, year));
    c2->SaveAs(Form("../plots/%s/sWeight/%s/BmassDataFit_sWeight_BsToJpsif2p_%s.pdf", year, mode, year));
   
    RooMsgService::instance().setSilentMode(true);
    RooStats::SPlot *sData = new RooStats::SPlot("sData", "An SPlot", *redData, model,
                                                 RooArgList(*nsigD, *nbkgE));

    const char* outpath=Form("../../red_ntuples/%s", year);
    TFile *outFile = new TFile(Form("%s/red_%s_Data_cutopt_sWeight_bdtr.root",
                               outpath, mode), "RECREATE");
    outFile->cd();

    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");

    std::cout << "Check SWeights:" << std::endl;
    std::cout << std::endl
                  << "N_{sig} " << nsigD->getVal() << ": From s_weight N_{sig} "
                  << sData->GetYieldFromSWeight(Form("nsigD_%s_%s", mode, year)) <<std::endl;

    std::cout << std::endl
                  << "N_{bkg} " << nbkgE->getVal() << ": From s_weight N_{bkg} "
                  << sData->GetYieldFromSWeight(Form("nbkgE_%s_%s", mode, year)) <<std::endl;

    std::cout << "Number of events = " << tr1->GetEntries() << std::endl;
    int flag=0;
    double sumSigSW = 0;
    double sumBkgSW = 0;
    for (Int_t i = 0; i < tr1->GetEntries(); i++){
         
        if (i%5000 == 0) std::cout << "--- ... Processing event: " << i << std::endl;
        s_wt = sData->GetSWeight(i, Form("nsigD_%s_%s", mode, year));
        
        if(s_wt<0){ // count number of -ive sWeights
            flag=flag+1;
        }
        b_wt = sData->GetSWeight(i, Form("nbkgE_%s_%s", mode, year));

        sumSigSW = sumSigSW+s_wt;
        sumBkgSW = sumBkgSW+b_wt;
        
        _s_wt->Fill();
        _b_wt->Fill();

    }

    std::cout << "Sum of signal sWeights: " << sumSigSW << std::endl;
    std::cout << "Sum of background sWeights: " << sumBkgSW << std::endl;

    std::cout << "Loop Finished" << std::endl;
    std::cout << "Number of events with -ive weights: " << flag << endl;

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


void getSwt(const char* mode, const char* year)
{
	gROOT->SetBatch(1);
	const char* file = Form("../../red_ntuples/%s", year);
    const char* filename = Form("red_%s_Data_cutopt_bdtr.root", mode);

	TChain* ch = new TChain("tree");
	ch->Add(Form("%s/%s", file, filename));
	TTree* tr = ch;
	int entries = tr->GetEntries();
	cout << "\n #entries-> " << entries << endl;

    TFile* fin = new TFile(Form("../fits/wspaces/ws_%s_DataFitResults_%s.root", mode, year));
    RooWorkspace* w3 = (RooWorkspace*)fin->Get("w3");

    RooFitResult* fitres=(RooFitResult*)w3->genobj(Form("fitresult_model_%s_%s_data", mode, year));
    RooDataSet *redData =(RooDataSet*)w3->data("data");
    RooAbsArg* nsigDarg = fitres->floatParsFinal().find(Form("nsigD_%s_%s", mode, year));
    RooRealVar* nsigD = dynamic_cast<RooRealVar*>(nsigDarg);
    RooAbsArg* nbkgEarg = fitres->floatParsFinal().find(Form("nbkgE_%s_%s", mode, year));
    RooRealVar* nbkgE = dynamic_cast<RooRealVar*>(nbkgEarg);

	RooAbsPdf* model = w3->pdf(Form("model_%s_%s", mode, year));
    std::cout << "Integral of model: " << model->getNorm() << std::endl;

    fitres->Print("v");
	
	TCanvas *c3 = new TCanvas("c3", "c3", 950, 1000);
	gPad->SetTickx();
	gPad->SetTicky();
    RooPlot *xframe2 = Bmass.frame(Title(""), Bins(50));
    xframe2->SetTitle("");
    xframe2->GetXaxis()->SetTitle("#bf{m(K^{+}K^{-}#mu^{+}#mu^{-}) [GeV]}");
    xframe2->GetYaxis()->SetTitleOffset(1.4);

	redData->plotOn(xframe2,RooFit::Name("Data"), MarkerStyle(8));
	model->plotOn(xframe2, RooFit::Name("Full PDF_data"), LineColor(kBlue), LineWidth(3));
	RooHist* hpull2 = xframe2->pullHist();
	
	model->plotOn(xframe2, 
                  RooFit::Name("CBall_data"), 
                  Components(Form("CB_%s_%s", mode, year)), 
                  LineColor(kGreen), 
                  LineWidth(3));
    model->plotOn(xframe2, 
                  RooFit::Name("bkgE_data"), 
                  Components(Form("bkgE_%s_%s", mode, year)), 
                  LineColor(kRed), 
                  DrawOption("F"), 
                  FillColor(kRed), 
                  FillStyle(3003));
    model->plotOn(xframe2, 
                  RooFit::Name("cb"), 
                  Components(Form("cb_%s_%s", mode, year)), 
                  LineColor(kMagenta), 
                  DrawOption("F"), 
                  FillColor(kMagenta), 
                  FillStyle(3008));

	xframe2->Draw();

    c3->SaveAs(Form("../plots/%s/sWeight/%s/BmassDataFit_sWeight_%s_%s.pdf", year, mode, mode, year));
    c3->SaveAs(Form("../plots/%s/sWeight/%s/BmassDataFit_sWeight_%s_%s.png", year, mode, mode, year));

	
	//calculating s_wt
	RooMsgService::instance().setSilentMode(true);

	RooStats::SPlot *sData = new RooStats::SPlot("sData", "An SPlot", *redData, model, 
                                                 RooArgList(*nsigD, *nbkgE));

	const char* outpath=Form("../../red_ntuples/%s", year);
	TFile *outFile = new TFile(Form("%s/red_%s_Data_cutopt_sWeight_bdtr.root", 
                               outpath, mode), "RECREATE");
	outFile->cd();
	
	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");

	std::cout << "Check SWeights:" << std::endl;
	
	std::cout << std::endl
                  << "N_{sig} " << nsigD->getVal() << ": From s_weight N_{sig} "
                  << sData->GetYieldFromSWeight(Form("nsigD_%s_%s", mode, year)) <<std::endl;

        std::cout << std::endl
                  << "N_{bkg} " << nbkgE->getVal() << ": From s_weight N_{bkg} "
                  << sData->GetYieldFromSWeight(Form("nbkgE_%s_%s", mode, year)) <<std::endl;

        std::cout << "Number of events = " << tr1->GetEntries() << std::endl;

    double sumSigSW = 0;
    double sumBkgSW = 0;
    int flag=0;

	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, Form("nsigD_%s_%s", mode, year));
        b_wt = sData->GetSWeight(i, Form("nbkgE_%s_%s", mode, year));

        if(s_wt<0) flag=flag+1;

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

        sumSigSW = sumSigSW+s_wt;
        sumBkgSW = sumBkgSW+b_wt;

        }

    std::cout << "Sum of signal sWeights: " << sumSigSW << std::endl;
    std::cout << "Sum of background sWeights: " << sumBkgSW << std::endl;
    std::cout << "Events with -ive weights: " << flag << endl;
	std::cout << "Loop Finished" << std::endl;

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

void splot(const int index, int yr)
{

    ostringstream str1;
    str1 << yr;
    string y=str1.str();
    const char* year=y.c_str();

    const int len=2;
    const char* mode[len] = {"BsToJpsiPhi",
                             "BsToJpsif2p"
    };

    if(strcmp(mode[index],"BsToJpsiPhi")==0){
        getSwt(mode[index], year);
    }
    else{
        cout << "Running on BsToJpsif2p" << endl;
        sWeightsBsToJpsif2p(mode[index], year);
    }
}