#include "Style/tdrstyle.C"
#include "Style/CMS_lumi.C"

#include "TCanvas.h"
#include "TPad.h"
#include "TFrame.h"
#include "TAxis.h"
#include "TLegend.h"
#include "TSystem.h"
#include <string>
#include <algorithm>


float mass_l = 8.0;
float mass_h = 14.0;
int nbins = 60;

using namespace RooFit;

void ppFit(){

setTDRStyle();
writeExtraText= false;
const char* filename2 = "Integrated/FIT_PbPb_DATA_sig_2_Bkg_ErfExp1_262620_263757_0cent200_0p0pt30p0_0p0y2p4.root";
const char* filename1 = "Integrated/FIT_PP_DATA_sig_2_Bkg_ErfExp1_262157_262328_-1cent1_0p0pt30p0_0p0y2p4.root";

TFile *f1 = new TFile(filename1);
RooWorkspace *_ws1 = (RooWorkspace*) f1->Get("myws"); 
_ws1->Print("v");
//opening file 2
TFile *f2 = new TFile(filename2);
RooWorkspace *_ws2 = (RooWorkspace*) f2->Get("myws");

TCanvas* comb = new TCanvas("comb", "comb",800,800);
//comb->SetFillStyle(4000);
//comb->SetFrameFillStyle(4000);


comb->cd();

//TPad *pad1 = new TPad("pad1","",0,0.228,1,1);
//TPad *pad2 = new TPad("pad2","",0,0,1,0.225);
TPad *pad1 = new TPad("pad1","",0,0,1,1);
TPad *pad2 = new TPad("pad2","",0,0,1,0);
pad1->SetFillStyle(4000); //will be transparent
pad1->SetFrameFillStyle(4000);
     pad2->SetTopMargin(0.02);
     pad2->SetBottomMargin(0.4);
     pad2->SetFillStyle(4000);
     pad2->SetFrameFillStyle(4000);
     pad2->SetFillColor(0);
     //pad1->SetBottomMargin(0.015);
     pad1->SetFillStyle(0);

pad1->Draw();
//pad2->Draw();
pad1->cd();

RooRealVar* mass =(RooRealVar*) _ws2->var("invariantMass"); //
RooDataSet* data0_fit =(RooDataSet*) _ws1->data("dataOS");
cout<<"pbpb data"<<data0_fit->numEntries()<<endl;
RooAbsPdf* pdf = (RooAbsPdf*) _ws1->pdf("pdf");  

RooRealVar* mass2 =(RooRealVar*) _ws2->var("invariantMass"); //
RooDataSet* data0_fit2 =(RooDataSet*) _ws2->data("dataOS");
cout<<"pp data"<<data0_fit2->numEntries()<<endl;
RooAbsPdf* pdf2 = (RooAbsPdf*) _ws2->pdf("pdf");  


RooPlot* frame_overlay = mass->frame(Bins(nbins),Range(mass_l,mass_h));
RooPlot* frame2 = mass->frame(Bins(nbins),Range(mass_l,mass_h));
frame_overlay->SetFillStyle(4000); //will be transparent

// cout<< " ok, " << frame_overlay->GetFillStyle() << endl ; //will be transparent
//  frame_overlay->SetFrameFillStyle(4000);

//defined from existing workspaces
// 1S yields
RooRealVar *nsig1f =(RooRealVar*) _ws1->var("N_{#varUpsilon(1S)}");
cout <<"Y(1S) yield (HI) : = "<<nsig1f->getVal()<<" +/- "<<nsig1f->getError()<<endl<<endl;
RooRealVar *nsig1f_pp =(RooRealVar*) _ws2->var("N_{#varUpsilon(1S)}");
cout <<"Y(1S) yield (pp) : = "<<nsig1f_pp->getVal()<<" +/- "<<nsig1f_pp->getError()<<endl<<endl;

//2S yields
RooRealVar *nsig2f   =(RooRealVar*) _ws1->var("R_{#frac{2S}{1S}}");
cout <<"Y(2S) yield (HI) : = "<<nsig2f->getVal()<<" +/- "<<nsig2f->getError()<<endl<<endl;
RooRealVar *nsig2f_pp =(RooRealVar*) _ws2->var("R_{#frac{2S}{1S}}");
cout <<"Y(2S) yield (pp) : = "<<nsig2f_pp->getVal()<<" +/- "<<nsig2f_pp->getError()<<endl<<endl;

//3S yields
RooRealVar *nsig3f   =(RooRealVar*) _ws1->var("R_{#frac{3S}{1S}}");
cout <<"Y(3S) yield (HI) : = "<<nsig3f->getVal()<<" +/- "<<nsig3f->getError()<<endl<<endl;
RooRealVar *nsig3f_pp =(RooRealVar*) _ws2->var("R_{#frac{3S}{1S}}");
cout <<"Y(3S) yield (pp) : = "<<nsig3f_pp->getVal()<<" +/- "<<nsig3f_pp->getError()<<endl<<endl;

// some additional definitions that we may not all need in the end.
RooRealVar *nbkgd   =(RooRealVar*) _ws1->var("n_{Bkgd}"); 
RooRealVar *f2Svs1S       = (RooRealVar*) _ws1->var("R_{#frac{2S}{1S}}");
RooRealVar *f3Svs1S       = (RooRealVar*) _ws1->var("R_{#frac{3S}{1S}}");
RooRealVar *width = (RooRealVar*) _ws1->var("width");
RooRealVar *turnOn =(RooRealVar*) _ws1->var("turnOn");
RooRealVar *decay =(RooRealVar*) _ws1->var("decay");

RooRealVar *nbkgd_pp   =(RooRealVar*) _ws2->var("n_{Bkgd}"); 
RooRealVar *f2Svs1S_pp       = (RooRealVar*) _ws2->var("R_{#frac{2S}{1S}}");
RooRealVar *f3Svs1S_pp       = (RooRealVar*) _ws2->var("R_{#frac{3S}{1S}}");
RooRealVar *width_pp = (RooRealVar*) _ws2->var("width");
RooRealVar *turnOn_pp =(RooRealVar*) _ws2->var("turnOn");
RooRealVar *decay_pp =(RooRealVar*) _ws2->var("decay");


// pdf definitions
RooAbsPdf *sig1S = ( RooAbsPdf*) _ws1->var("sig1S");
RooAbsPdf *sig2S = ( RooAbsPdf*) _ws1->var("sig2S");
RooAbsPdf *sig3S = ( RooAbsPdf*) _ws1->var("sig3S");
RooAbsPdf *pdf_combinedbkgd = ( RooAbsPdf*) _ws1->var("bkgPdf");
RooAbsPdf *pdf_combinedbkgd_pp = ( RooAbsPdf*) _ws2->var("bkgPdf");

//Then, I plot the pbpb data and its fit, and I set the pp bkgd parameters to the pbpb ones :

data0_fit->plotOn(frame_overlay,Name("dataOS_FIT"),MarkerSize(1.2));
pdf->plotOn(frame_overlay,Name("thePdf"));
RooHist *hPull = frame_overlay->pullHist(0,0,true);

hPull->SetName("hPull");
hPull->SetMarkerSize(1.2);
frame2 -> addPlotable(hPull,"PX");
pdf->plotOn(frame_overlay,Name("sig1S"),Components("sig1S"),LineColor(kGray+2),LineStyle(3),LineWidth(3));
pdf->plotOn(frame_overlay,Components("sig2S"),LineColor(kGray+2),LineStyle(3),LineWidth(3));
pdf->plotOn(frame_overlay,Components("sig3S"),LineColor(kGray+2),LineStyle(3),LineWidth(3));
pdf->plotOn(frame_overlay,Name("bkgPdf"),Components("bkgPdf"),LineStyle(4),LineColor(kBlue),LineWidth(3));
// try fixing these.

width_pp->setVal(width->getVal());
turnOn_pp->setVal(turnOn->getVal());
decay_pp->setVal(decay->getVal());


// Now the important part:
//1. make a new workspace like you would do a simultaneous fit (but you dont need to fit since this was done already)

RooWorkspace* ws_sim = new RooWorkspace ("ws_sim","ws for simfit");
ws_sim->var("invariantMass");
RooCategory dataCat("dataCat", "dataCat");
dataCat.defineType("hi");
dataCat.defineType("pp");
// 2. Define the simultaneous pdf:

RooDataSet data("data", "data", *mass, Index(dataCat), Import("hi", *data0_fit), Import("pp", *data0_fit2));
data.Print("v");
ws_sim->import(data);
RooSimultaneous simPdf("simPdf", "simPdf", dataCat); 
// RooAbsPdf * pdf_hi = ws_sim->pdf("pdf_hi");
// RooAbsPdf * pdf_pp = ws_sim->pdf("pdf_pp");
assert(pdf);
assert(pdf2);
simPdf.addPdf(*pdf, "hi");
simPdf.addPdf(*pdf2, "pp");
ws_sim->import(simPdf,Silence(),
	 // RooFit::RenameAllNodes((hi)?"hi":"pp"),
	 // RooFit::RenameAllVariablesExcept((hi)? "hi": "pp", 
	 // 				  "npow,invariantMass,"
	 // 				  "nsig1_pp,nsig2_pp,nsig3_pp,"
	 // 				  "raa1S,raa2S,raa3S," 
	 // 				  "mean,"
	 // 				  "turnOn,"
	 // 				  "alpha,"
	 // 				  "sigma1"
	 // 				  ),
	 RooFit::RecycleConflictNodes());

//3. Plotting PbPb with normalisation = 1:
float Nexp = 1; //simPdf.expectedEvents(data.get());

data.plotOn(frame_overlay,Invisible(),
	Cut("(dataCat==dataCat::hi)"),
	Name("data_norm"));

simPdf.plotOn(frame_overlay, Slice(dataCat,"hi"),
	Normalization(Nexp),
	ProjWData(dataCat,data),
	Name("full_FIT"),
	LineColor(kBlue),
	LineWidth(3)
	);
/*simPdf.plotOn(frame_overlay, Slice(dataCat,"hi"),
  		Normalization(Nexp),
  		Components("bkgPdf"),
  		ProjWData(dataCat,data),
  		Name("Hi_FIT"),
  		//	LineColor(kRed),
  		LineWidth(6)
  		); 

*/
/*
float bkg = nbkgd -> getVal();
cout<<"bgsfdaf"<<bkg<<endl;
nbkgd_pp->setVal(104225); // ~=nbkgd_pbpb
nsig1f_pp->setVal(nsig1f_pp->getVal()*(nsig1f->getVal()/nsig1f_pp->getVal()));
nsig2f_pp->setVal(nsig2f_pp->getVal()*(nsig1f->getVal()/nsig1f_pp->getVal()));
nsig3f_pp->setVal(nsig3f_pp->getVal()*(nsig1f->getVal()/nsig1f_pp->getVal()));
simPdf.plotOn(frame_overlay, Slice(dataCat,"pp"),
  		  Normalization(0.65),
  		  ProjWData(dataCat,data),
  		  Name("pp_FIT"),
  		  LineColor(kRed),
  		  LineWidth(4),
  		  LineStyle(7)
  		  );
*/
//frame_overlay->getAttText()->SetTextSize(0.022);
//frame_overlay->getAttFill()->SetFillStyle(0);
frame_overlay->SetTitle("");
//frame_overlay->GetXaxis()->SetTitle("");
frame_overlay->GetXaxis()->CenterTitle(kTRUE);

//frame_overlay->GetYaxis()->SetTitleFont(43);

//frame_overlay->GetYaxis()->SetTitleSize(28);
//frame_overlay->GetYaxis()->SetLabelFont(43);
//frame_overlay->GetYaxis()->SetLabelSize(28);
//frame_overlay->GetYaxis()->SetTitleOffset(2.0);
TGaxis::SetMaxDigits(4);

//frame_overlay->GetXaxis()->SetTitleFont(43);
frame_overlay->GetXaxis()->SetTitleOffset(1.10);
frame_overlay->GetXaxis()->CenterTitle();
//frame_overlay->GetXaxis()->SetTitleSize(28);
//frame_overlay->GetXaxis()->SetLabelFont(43);
//frame_overlay->GetXaxis()->SetLabelSize(28);
frame_overlay->GetXaxis()->SetTitle("m_{#mu#mu} (GeV/c^{2})");

frame_overlay->Draw();

TLatex *t = new TLatex(); t->SetNDC();
t->SetTextSize(28);
t->SetTextFont(43);
float dy = 0.04; 

t->DrawLatex(0.42, 0.87-dy, Form("p^{#mu#mu}_{T} < %.0f GeV/c",30.)); dy+=0.06;
t->DrawLatex(0.42, 0.87-dy, Form("|y^{#mu#mu}| < %.1f",2.4)); dy+=0.06;
t->DrawLatex(0.42, 0.87-dy, Form("p_{T}^{#mu} > %.0f GeV/c",4.0)); dy+=0.06;

TLegend* leg = new TLegend(0.7,0.48,0.9,0.7); leg->SetTextSize(28);
leg->SetTextFont(43);
leg->SetHeader("pp");
leg->AddEntry(frame_overlay->findObject("dataOS_FIT"),"Data","pe");
leg->AddEntry(frame_overlay->findObject("full_FIT"),"Total fit","l");
leg->AddEntry(frame_overlay->findObject("sig1S"),"Signal","l");
leg->AddEntry(frame_overlay->findObject("bkgPdf"),"Background","l");

leg->Draw("same");

TString label;
label = "";

CMS_lumi(pad1, 104, 33, label);


pad1->Update();
pad1->SetBottomMargin(0.15);
comb->Update();

//pad2->cd();
   TLine * pline = new TLine(mass_l,0.0,mass_h,0.0);


     frame2->SetFillColor(0);
     frame2->GetYaxis()->CenterTitle(kTRUE);
     frame2->GetYaxis()->SetTitleOffset(1.5);
     frame2->GetYaxis()->SetTitleFont(43);
     frame2->GetYaxis()->SetLabelFont(43);
     frame2->GetYaxis()->SetTitleSize(20);
     frame2->GetYaxis()->SetLabelSize(15);
     frame2->GetYaxis()->SetTitle("Pull");
     frame2->GetXaxis()->CenterTitle(kTRUE);
     frame2->GetYaxis()->SetRangeUser(-4.5,4.5);
     frame2->GetYaxis()->SetNdivisions(508);

     frame2->GetXaxis()->SetTitleFont(43);
     frame2->GetXaxis()->SetTitleOffset(4.);
     frame2->GetXaxis()->SetTitleSize(28);
     frame2->GetXaxis()->SetLabelFont(43);
     frame2->GetXaxis()->SetLabelSize(28);
     frame2->GetXaxis()->SetTitle("m_{#mu#mu} (GeV/c^{2})");

//frame2->Draw();
//pline->Draw("same");
//pad2->Update();
pad1->Update();


comb->SaveAs("FIT_PP_05042017.pdf");
comb->SaveAs("FIT_PP_05042017.png");

}