#include "TSystem.h"
#include "TChain.h"
#include "TFile.h"
#include "TTimeStamp.h"
#include "TH1F.h"
#include "TMath.h"
#include <iostream>
#include <fstream>
#include "TH2F.h"
#include<math.h>
#include <TProfile.h>
#include <TColor.h>
#include "THashList.h"
#include "TCanvas.h"
#include "TClonesArray.h"
#include "TTree.h"
#include "RooAbsPdf.h"
#include "RooAbsPdf.h"
#include "RooAbsData.h"

void myLegendSetUp(TLegend *currentLegend,float currentTextSize,int columns);



void Fit(Float_t centMin=50.0, Float_t centMax=70.0, Double_t minMass=1.68, Double_t maxMass=5, Double_t ptMin=0.0, Double_t ptMax=10.0) {


gSystem->Load("libRooFit");
using namespace RooFit ;

  TFile *fData_file = TFile::Open(Form("data/DataTreeSE_cent_%.0f_%.0f.root",centMin,centMax),"read");
  TFile* fComponents= TFile::Open("fComponents.root");

  Double_t fPt, fMass;

  TTree *fTreeDataPtMass = new TTree("fTreeDataPtMass", "fTreeDataPtMass");
  fTreeDataPtMass->Branch("rPt", &fPt, "rPt/D");
  fTreeDataPtMass->Branch("rMass", &fMass, "rMass/D");

  TTree* fData_Tree  =(TTree*) fData_file->Get("fTreeDataPt");
  fData_Tree->SetBranchAddress("fPt", &fPt);
  fData_Tree->SetBranchAddress("fMass", &fMass);
  fData_Tree->Print();

  const Int_t nEntries = fData_Tree->GetEntries();


  for(Int_t ie=0; ie<nEntries;ie++){
    fData_Tree->GetEntry(ie);
    fTreeDataPtMass->Fill();

 }
//**********get bkg,signal and signal+bkg histograms ****************

  TH2D* hBkg=(TH2D*)fComponents->Get("hBkg");
  TH2D* hSplusB=(TH2D*)fComponents->Get("hSplusB");

 //**********get MC distributions*******************
  TH2D* hHadronic=(TH2D*)fComponents->Get(TString::Format("hInclusive_PtMass_cent_%.0f_%.0f",centMin, centMax));
  TH2D *hCoherent_Pt = (TH2D*)fComponents->Get(TString::Format("hCoherentJpsi_PtMass_cent_%.0f_%.0f",centMin, centMax));
 
  hBkg->SetLineColor(kMagenta);
  hCoherent_Pt->SetLineColor(kBlue);
  hHadronic->SetLineColor(kCyan);

//=============================RooFit=================================
  Int_t order = 0; // order of the interpolation between bins. Zero is no interpolation

  RooRealVar rPt("rPt","#it{p}_{T} (GeV/#it{c})",ptMin, ptMax);
  RooRealVar rMass("rMass","#m_{ee} (GeV/#it{c^{2}})",minMass,maxMass);

  const Int_t binsPt=(1.0-ptMin)/0.02;
  RooBinning ptBins(ptMin, 10.0);
  ptBins.addUniform(binsPt, ptMin, 1.0);
  ptBins.addUniform(9, 1.0, 10.0);

  const Int_t binsMass=(maxMass-minMass)/0.04;
  RooBinning MassBins(minMass, maxMass);
  MassBins.addUniform(binsMass, minMass ,maxMass);
 

//***************unbinned data***************************
  RooDataSet rData_UnBin("rData_UnBin","rData_UnBin",RooArgSet(rMass,rPt),Import(*fTreeDataPtMass));
//***************binned data***********************************
  RooDataHist rDataBinned("rDataBinned","rDataBinned",RooArgList(rMass,rPt),hSplusB);
  RooDataHist rDataBkg("rDataBkg","rDataBkg",RooArgList(rMass,rPt),hBkg);

//********define the cuts for mass and pT********************
  rPt.setRange("PtCut",2,3);
  rMass.setRange("MassWindow",2.92,3.16);

  RooRealVar r_Frc("r_Frc","Number of Jpsi",0,nEntries); 
  RooRealVar rME_Frc("rME_Frc","Number bkg",0,nEntries); //test with free parameter for the bkg component
  RooRealVar rHad_Frc("rHad_Frc","Number of hadronic Jpsi",0,10000); 

//********get the number of entries from the ME histogram************************
 
  Double_t Bkg_Frc=hBkg->Integral(hBkg->GetXaxis()->FindBin(minMass+0.001),hBkg->GetXaxis()->FindBin(maxMass-0.001),hBkg->GetYaxis()->FindBin(ptMin+0.001),hBkg->GetYaxis()->FindBin(ptMax-0.001));


  RooDataHist rBkg_hist("rBkg_hist","rBkg_hist",RooArgList(rMass,rPt),hBkg);
  RooDataHist rCoherent_hist("rCoherent_hist","rCoherent_hist",RooArgList(rMass,rPt),hCoherent_Pt);
  RooDataHist rHad_hist("rHad_hist","rHad_hist",RooArgList(rMass,rPt),hHadronic);

  RooHistPdf rBkg_Pdf("rBkg_Pdf","rBkg_Pdf",RooArgSet(rMass,rPt),rBkg_hist,order);
  RooHistPdf rCoherent_Pdf("rCoherent_Pdf","rCoherent_Pdf",RooArgSet(rMass,rPt),rCoherent_hist,order);
  RooHistPdf rHadr_Pdf("rHadr_Pdf","rHadr_Pdf",RooArgSet(rMass,rPt),rHad_hist,0);

  rCoherent_Pdf.forceNumInt(kTRUE);
  rHadr_Pdf.forceNumInt(kTRUE);
  rBkg_Pdf.forceNumInt(kTRUE);

  RooGenericPdf rCoherent_Frc("rCoherent_Frc","Number of coherent","r_Frc*0.597201",RooArgSet(r_Frc));

// // Create the model as the sum of templates
  
  RooAddPdf Model("Model","Extended sum templates",RooArgList(rBkg_Pdf,rCoherent_Pdf,rHadr_Pdf),
						 RooArgList( RooConst(Bkg_Frc),rCoherent_Frc,rHad_Frc));
 

  RooFitResult* r = Model.fitTo(rDataBinned,Save(),Extended(kTRUE));

  TLegend *myLegend1 = new TLegend(0.37,0.29,0.84,0.72);
  myLegendSetUp(myLegend1,0.04,1);
  
  myLegend1->AddEntry((TObject*)0,TString::Format("%.1f < M_{ ee} (GeV/#it{c}^{2}) < %.1f",2.92,3.16),"");
  myLegend1->AddEntry(hBkg," Background","l");
  myLegend1->AddEntry(hCoherent_Pt," Signal C","l");
  myLegend1->AddEntry(hHadronic,"Hadronic component","l");

 
  
  TCanvas *cRooFit = new TCanvas("cRooFit","cRooFit",1600,1100);
  cRooFit->Divide(2,1);
  cRooFit->cd(1);
   
  RooPlot* framePt = rPt.frame(Binning(ptBins),Title(" ")); 
  
  rDataBinned.plotOn(framePt,Binning(ptBins),CutRange("MassWindow")) ;
  Model.plotOn(framePt,Binning(ptBins),ProjectionRange("MassWindow"), LineColor(kBlack),Name("Model"));
  Model.plotOn(framePt,Components(rBkg_Pdf), LineColor(kMagenta),LineWidth(2),Binning(ptBins),ProjectionRange("MassWindow"),Name("Background")); 
  Model.plotOn(framePt,Components(rCoherent_Pdf), LineColor(kBlue),LineWidth(2),Binning(ptBins),ProjectionRange("MassWindow"),Name("Coherent Jpsi"));
  Model.plotOn(framePt,Components(rHadr_Pdf), LineColor(kCyan),LineWidth(2),Binning(ptBins),ProjectionRange("MassWindow"), Name("Hadronic"));
  framePt->Draw();
//   
  cRooFit->cd(2);
  RooPlot* frame = rMass.frame(Binning(MassBins),Title(" "));
  
  rDataBinned.plotOn(frame,Binning(MassBins),CutRange("PtCut")) ;
  Model.plotOn(frame,Binning(MassBins),ProjectionRange("PtCut"),LineColor(kBlack));
  Model.plotOn(frame,Components(rBkg_Pdf), LineColor(kMagenta),LineWidth(2),Binning(MassBins),ProjectionRange("PtCut"));
  Model.plotOn(frame,Components(rCoherent_Pdf), LineColor(kBlue),LineWidth(2),Binning(MassBins),ProjectionRange("PtCut"));
  Model.plotOn(frame,Components(rHadr_Pdf), LineColor(kCyan),LineWidth(2),Binning(ptBins),ProjectionRange("PtCut"), Name("Hadronic"));
  frame->Draw();
  myLegend1->Draw();
  
  TCanvas *cRooFit_QA = new TCanvas("cRooFit_QA","cRooFit_QA",1600,1100);
  cRooFit_QA->Divide(2,1);
  cRooFit_QA->cd(1);
  
  RooPlot* frameQA_Pt=rPt.frame(Binning(ptBins),Title(" ")); 
  rDataBinned.plotOn(frameQA_Pt,Binning(ptBins),CutRange("MassWindow")) ;
  Model.plotOn(frameQA_Pt,Components(rBkg_Pdf), LineColor(kMagenta),LineWidth(2),Binning(ptBins),ProjectionRange("MassWindow"),Name("Background"));
  rBkg_hist.plotOn(frameQA_Pt,Binning(ptBins),CutRange("MassWindow"),MarkerColor(kRed)) ;
  
  frameQA_Pt->Draw();
  
  cRooFit_QA->cd(2);
  
  RooPlot* frameQA_Mass=rMass.frame(Binning(MassBins),Title(" "));
  rDataBinned.plotOn(frameQA_Mass,Binning(MassBins),CutRange("PtCut")) ;
  Model.plotOn(frameQA_Mass,Components(rBkg_Pdf), LineColor(kMagenta),LineWidth(2),Binning(MassBins),ProjectionRange("PtCut"),Name("Background"));
  rBkg_hist.plotOn(frameQA_Mass,Binning(MassBins),CutRange("PtCut"),MarkerColor(kRed)) ;
  
  frameQA_Mass->Draw();
  
//   
//   


  
}

void myLegendSetUp(TLegend *currentLegend,float currentTextSize,int columns){
  currentLegend->SetTextFont(42);
  currentLegend->SetBorderSize(0);
  currentLegend->SetFillStyle(0);
  currentLegend->SetFillColor(0);
  currentLegend->SetMargin(0.25);
  currentLegend->SetTextSize(currentTextSize);
  currentLegend->SetEntrySeparation(0.5);
  currentLegend->SetNColumns(columns);
  return;
}