///////////////////////////////////////////////////////////////////////////////////////////////////////////
//Contact: Felipe (felipe@if.ufrj.br)
//
//  - How to run:
//      .L FitG2CB_simple_MC.C++
//      FitG2CB_simple_MC();
//
/////////////////////////////////////////////////////////////////////////////////////////////////////////// 

//ROOT libraries
#include "TStyle.h"
#include "TH1D.h"
#include "TCanvas.h"
#include "TLatex.h"
#include "TLine.h"
#include "TFile.h"
#include "TMath.h"
#include "TF1.h"

//RooFit libraries
#include "RooRealVar.h"
#include "RooDataHist.h"
#include "RooFormulaVar.h"
#include "RooAbsPdf.h"
#include "RooCBShape.h"
#include "RooExponential.h"
#include "RooBernstein.h"
#include "RooGaussian.h"
#include "RooAddPdf.h"
#include "RooExtendPdf.h"
#include "RooFitResult.h"
#include "RooPlot.h"
#include "RooHist.h"
#include "RooGlobalFunc.h"
#include "RooProduct.h"

using namespace std;
using namespace RooFit;

int FitG2CB_simple_MC(bool sumw2 = true, TString minimizer = "Minuit"){ 
    TFile *file = new TFile("histofile_MC.root", "READ");
    TH1D* h = (TH1D*) file->Get("histo");

    Int_t lhcbFont        = 132;  // Old LHCb style: 62;
    gStyle->SetLabelFont(lhcbFont,"x");
    gStyle->SetLabelFont(lhcbFont,"y");
    gStyle->SetLabelFont(lhcbFont,"z");
    gStyle->SetTitleFont(lhcbFont);
    gStyle->SetTitleFont(lhcbFont,"x");
    gStyle->SetTitleFont(lhcbFont,"y");
    gStyle->SetTitleFont(lhcbFont,"z");	
    gStyle->SetOptTitle(1);

    Double_t massLo = 1910;
    Double_t massHi = 2030;
    Double_t massLoRed = 1930;
    Double_t massHiRed = 2010;
    Double_t massSigLo = 1950;
    Double_t massSigHi = 1990;
    Double_t meanmass = 1970;

    TString particles = "K^{-}K^{+}#pi^{+}"; 

    TString name = h->GetName();
    RooRealVar *mass = new RooRealVar("mass","Mass",massLo,massHi,"MeV/#it{c}^{2}");
    mass->setRange("signalRange",massSigLo,massSigHi);
    mass->setRange("fullRange",massLo,massHi);
    RooDataHist data("data"+name,"data"+name, *mass, Import(*h));

    double xent = h->Integral(0,h->GetNbinsX()+1);

    //---> Run2 MC
    double meanRsigma21= 1.75;
    double meanRsigma31= 1.602; 
    double meanFrac1= 0.568;
    double meanFrac2= 0.36;
    double meanOffset= 0.0;
    double meanAlfa1= 1.61;
    double meanAlfa2= -2.24;
    double meanN1= 1.61;
    double meanN2= 2.96;
    double meanSigmaG1 = 4.60;

    RooRealVar *Sigma_G1 = new RooRealVar("#sigma G","mass resolution",meanSigmaG1,2.,20.,"GeV/#it{c}^{2}");
    RooRealVar *meanoffset = new RooRealVar("#Delta M","",meanOffset,-10.0,10.0);
    RooRealVar *R_sigma21 = new RooRealVar("R #sigma 21","",meanRsigma21,1,5.5); 
    RooRealVar *R_sigma31 = new RooRealVar("R #sigma 31","",meanRsigma31,1,5.5); 
    RooRealVar *sig1frac = new RooRealVar("f_{G}","fraction of component 1 in signal",meanFrac1,0.0,1.0, "");
    RooRealVar *sig2frac = new RooRealVar("f_{CB2}","fraction of component 2 in signal",meanFrac2,0.0,0.5, "");
    RooRealVar *N1  = new RooRealVar("N1","CBsinal",meanN1,0.,10);
    RooRealVar *Alfa1  = new RooRealVar("Alfa1","CBsinal",meanAlfa1,0.0,10);
    RooRealVar *N2  = new RooRealVar("N2","CBsinal",meanN2,0.,10);
    RooRealVar *Alfa2  = new RooRealVar("Alfa2","CBsinal",meanAlfa2,-10, 0.0);

    RooRealVar *m = new RooRealVar("mG","Mean mass", meanmass, meanmass-5, meanmass+5,"GeV/c^{2}");
    RooFormulaVar *m3 = new RooFormulaVar("m_{G3}", "@0+@1", RooArgList(*m, *meanoffset));
    RooFormulaVar *Sigma_CB3 = new RooFormulaVar("#sigma_{CB3}"," CB sigma ","@0*@1" ,RooArgList(*Sigma_G1,*R_sigma31));
    RooFormulaVar *Sigma_CB2 = new RooFormulaVar("#sigma_{CB2}"," CB sigma ","@0*@1" ,RooArgList(*Sigma_G1,*R_sigma21));

    RooRealVar *N_Sig = new RooRealVar ("N_{sig}","Number of signal events",0.5*xent,0.,xent);	

    meanoffset->setConstant(true);

    //signal pdf
    RooAbsPdf* pdf_Sig1 = new RooGaussian("pdf_Sig1","D signal mass pdf",*mass,*m,*Sigma_G1);//media e sigma
    RooAbsPdf* pdf_Sig2 = new RooCBShape("pdf_Sig2","sinal",*mass,*m,*Sigma_CB2,*Alfa1,*N1);//media, sigma
    RooAbsPdf* pdf_Sig3 = new RooCBShape("pdf_Sig3","sinal",*mass,*m3,*Sigma_CB3,*Alfa2,*N2);//media, sigma
    RooAbsPdf *pdf_Sig = new RooAddPdf("pdf_Sig","Total GCB Signal  PDF",RooArgSet(*pdf_Sig1,*pdf_Sig2,*pdf_Sig3),RooArgList(*sig1frac,*sig2frac),kTRUE);
    RooAbsPdf* epdf_Sig = new RooExtendPdf("epdf_Sig","extended signal p.d.f",*pdf_Sig,*N_Sig,"fullRange");

    // Fit model for signal (add bkg model when needed)
    RooAbsPdf* pdf_fit = new RooAddPdf("pdf_fit","Total(B+S) D(s) -> K K #pi PDF",*epdf_Sig); 

    //Perform fit of model to corresponding data 
    RooFitResult *fit = pdf_fit->fitTo(data,Extended(true),Strategy(1),SumW2Error(sumw2),Minimizer(minimizer),Save(true)); 

    RooPlot* frame1 = mass->frame(Title(" "));
    data.plotOn(frame1,Name("fit_"+name));
    pdf_fit->plotOn(frame1,Name("fitPDF_"+name),LineWidth(3),LineColor(2));
    pdf_fit->paramOn(frame1,Layout(0.65, 0.95, 0.95),Format("NEU",AutoPrecision(1)));

    RooPlot* frame2 = mass->frame(Title(" "));
    data.plotOn(frame2,Name("fit"+name));
    pdf_fit->plotOn(frame2,Name("fitPDF_"+name),LineWidth(3),LineColor(2));
    pdf_fit->plotOn(frame2,Components("epdf_Sig"),LineWidth(2),LineStyle(kDashed),LineColor(4));
    pdf_fit->plotOn(frame2,Components("pdf_Sig1"),LineWidth(2),"",LineColor(kBlue-2));
    pdf_fit->plotOn(frame2,Components("pdf_Sig2"),LineWidth(2),"",LineColor(kBlue-4));
    pdf_fit->plotOn(frame2,Components("pdf_Sig3"),LineWidth(2),"",LineColor(kBlue-6));

    RooPlot *frame1pulls= mass->frame(Title(" "));
    frame1pulls->addPlotable((frame1->pullHist("fit_"+name, "fitPDF_"+name)), "P"); 
    frame1pulls->SetMinimum(-5); 
    frame1pulls->SetMaximum(+5); 

    int nParFit = fit->floatParsFinal().getSize();
    int nBinsPlot = frame1->GetNbinsX();
    int ndfFit = nBinsPlot - nParFit;	
    double chi2red = frame1->chiSquare("fitPDF_"+name, "fit_"+name);
    double prob=TMath::Prob(chi2red*nBinsPlot,nBinsPlot);
    double effSigma = sqrt(Sigma_G1->getVal()*Sigma_G1->getVal()*sig1frac->getVal()+(1-sig1frac->getVal())*sig2frac->getVal()*Sigma_CB2->getVal()*Sigma_CB2->getVal()+(1-sig1frac->getVal())*(1-sig2frac->getVal())*Sigma_CB3->getVal()*Sigma_CB3->getVal()); //sqrt[fG1*sigma1² + (1-fG1)*fG2*sigma2² + (1-fG1)(1-fG2)*sigma3²]
    
    std::cout << h->GetName() << ": chi2/ndf = " << chi2red << ", p-value = " << prob << ", Fit covQual() = " << fit->covQual() << std::endl;
    std::cout << "G2CBpol2 :: Fit results" << std::endl;

    std::cout << std::endl << std::endl;
    std::cout << "Sigma_G1 = " << Sigma_G1->getVal() << " Sigma_G1^2 = " << Sigma_G1->getVal()*Sigma_G1->getVal() << std::endl;
    std::cout << "Sigma_CB2 = " << Sigma_CB2->getVal() << " Sigma_CB2^2 = " << Sigma_CB2->getVal()*Sigma_CB2->getVal() << std::endl;
    std::cout << "Sigma_CB3 = " << Sigma_CB3->getVal() << " Sigma_CB3^2 = " << Sigma_CB3->getVal()*Sigma_CB3->getVal() << std::endl;
    std::cout << "Sig1frac = " << sig1frac->getVal() << " Sig2frac = " << sig2frac->getVal() << std::endl;
    std::cout << "(1-fG1)*fCB2 = " << (1-sig1frac->getVal())*sig2frac->getVal() << std::endl;
    std::cout << "(1-fG1)*(1-fCB2) = " << (1-sig1frac->getVal())*(1-sig2frac->getVal()) << std::endl;

    std::cout << "eff sigma = " << effSigma << std::endl << std::endl;

    Double_t Nsig = N_Sig->getVal();
    Double_t NsigError = N_Sig->getPropagatedError(*fit);

    TLatex *lhcbLatex = new TLatex(0.20, 0.80,"#splitline{LHCb}{#scale[0.9]{#sqrt{s} = 13 TeV }}");
    TLatex *lhcbLatex2 = new TLatex(0.20, 0.80,"#splitline{LHCb}{#scale[0.9]{#sqrt{s} = 13 TeV }}");
    lhcbLatex->SetNDC(kTRUE);
    lhcbLatex->SetTextColor(1);
    lhcbLatex->SetTextSize(0.05);
    lhcbLatex2->SetNDC(kTRUE);
    lhcbLatex2->SetTextColor(1);
    lhcbLatex2->SetTextSize(0.05);
    frame1->addObject(lhcbLatex);
    frame2->addObject(lhcbLatex2);

    int exp = TMath::Log(Nsig)/TMath::Log(1000);
    if (exp==0) exp=0;

    TString sufix = " kMGTPE";
    TString yield = Form("%.3g %c", Nsig/pow(1000, exp), sufix[exp]);
    TLatex *dataInfo = new TLatex(0.20, 0.60,"#scale[0.9]{Signal Yield = "+yield+"}");
    TLatex *dataInfo2 = new TLatex(0.20, 0.60,"#scale[0.9]{Signal Yield = "+yield+"}");
    dataInfo->SetNDC(kTRUE);
    dataInfo->SetTextColor(1);
    dataInfo->SetTextSize(0.05);
    dataInfo2->SetNDC(kTRUE);
    dataInfo2->SetTextColor(1);
    dataInfo2->SetTextSize(0.05);
    frame1->addObject(dataInfo);
    frame2->addObject(dataInfo2);

    gStyle->SetOptTitle(1);

    TCanvas* c1 = new TCanvas("c1"+name,"fitpdf",1200,800) ;
    gPad->SetLeftMargin(0.15) ;
    TPad *pad1 = new TPad("pad1","top pad",0,0.32,.98,.98);
    pad1->Draw();
    TPad *pad2 = new TPad("pad2","bottom pad",0,0,.98,0.30);
    pad2->Draw();
    pad1->cd();
    gPad->SetTopMargin(0.12);
    gPad->SetLeftMargin(0.16);
    gPad->SetRightMargin(0.04);
    gPad->SetBottomMargin(0.2);
    frame1->GetYaxis()->SetTitle("Candidates / [MeV/#it{c}^{2}]");
    frame1->GetXaxis()->SetTitleSize(.07); 
    frame1->GetYaxis()->SetTitleSize(.07); 
    frame1->GetYaxis()->SetLabelSize(.05); 
    frame1->GetYaxis()->SetNdivisions(505);
    frame1->GetXaxis()->SetLabelSize(.07); 
    //  frame1->getAttText()->SetTextSize(0.070); 
    frame1->SetTitle(Form("chi2/ndf=%.2f p-value=%.2f covQual=%d", chi2red,prob,fit->covQual()));
    frame1->SetTitleOffset(1.05);
    frame1->SetXTitle("Mass ( "+particles+" ) [MeV/#it{c}^{2}]");
    frame1->Draw();
    pad2->cd();
    gPad->SetLeftMargin(0.16);
	gPad->SetRightMargin(0.04);
	gPad->SetTopMargin(0.01);
	gPad->SetBottomMargin(0.25);
	frame1pulls->SetTitle("");
	frame1pulls->GetYaxis()->SetTitle("Pulls"); 
	frame1pulls->GetYaxis()->SetNdivisions(505);
	frame1pulls->GetYaxis()->SetLabelSize(.11); 
	frame1pulls->GetYaxis()->SetTitleSize(.14); 
	frame1pulls->GetXaxis()->SetLabelSize(.12); 
	frame1pulls->GetXaxis()->SetTitleSize(.15); 
	frame1pulls->GetXaxis()->SetTickLength(.07); 
	frame1pulls->GetYaxis()->SetTitleOffset(0.45); 
	frame1pulls->GetXaxis()->SetTitleOffset(0.75); 
	frame1pulls->SetXTitle("Mass ( "+particles+" ) [MeV/#it{c}^{2}]");
	TLine *l0= new TLine(massLo, 0, massHi, 0.);
	TLine *lm3= new TLine(massLo, -3, massHi, -3);
	TLine *lp3= new TLine(massLo, 3, massHi, 3);
	l0->SetLineStyle(1);l0->SetLineColor(2);
	lm3->SetLineStyle(2);lm3->SetLineColor(2);
	lp3->SetLineStyle(2);lp3->SetLineColor(2);
	frame1pulls->addObject(l0);
	frame1pulls->addObject(lm3);
	frame1pulls->addObject(lp3);
	frame1pulls->Draw();

    // c1->Print("../fits/Fit_G2CB.png","png"); 	
	
	gStyle->SetOptTitle(0);
	TCanvas* c2 = new TCanvas("c2"+name,"fitpdf",1200,800) ;
	gPad->SetLeftMargin(0.15) ;
	gPad->SetTopMargin(0.12);
	gPad->SetLeftMargin(0.16);
	gPad->SetRightMargin(0.04);
	gPad->SetBottomMargin(0.2);
	frame2->GetYaxis()->SetTitle("Candidates / [MeV/#it{c}^{2}]");
	frame2->GetXaxis()->SetTitleSize(.07); 
	frame2->GetYaxis()->SetTitleSize(.07); 
	frame2->GetYaxis()->SetLabelSize(.05); 
	frame2->GetYaxis()->SetNdivisions(505);
	frame2->GetXaxis()->SetLabelSize(.07); 
	frame2->SetTitleOffset(1.05);
	frame2->SetXTitle("Mass ( "+particles+" ) [MeV/#it{c}^{2}]");
	frame2->Draw();

    // c2->Print("../fits/Fit_G2CB_without_parameters.png","png"); 

    std::cout << "\n==> Nsig = " << Nsig << " +- " << NsigError << ", covQual() = " << fit->covQual() << "\n" << std::endl;

    return fit->covQual();
}