#include "Riostream.h"
#include "TFile.h"
#include "TTree.h"
#include "TCanvas.h"
#include "TH1D.h"
#include "TF1.h"
#include "TLegend.h"
#include "TStyle.h"
#include "TMatrixDSym.h"
#include "TFitResult.h"
#include "TGraph.h"
#include "TMultiGraph.h"

using namespace std;

void fitDeltaEFinal(){
    
    const double min_de = -0.15;
    const double max_de =  0.15;
    //define an histogram to look at deltaE distribution
    TH1D* h_data = new TH1D("h_data",";#DeltaE [GeV]; Entries",40,min_de,max_de);

    //open file and take the tree
    TFile* file = TFile::Open("data.root");
    TTree* tree = (TTree*) file->Get("treeData");
    
    int tot_entries = tree->GetEntries();
    cout << "Total entries in the tree: " << tot_entries << endl;
    
    //link the variables with tree banches
    double B_de;
    double bkg_killer;
    tree->SetBranchAddress("B_de",&B_de);
    tree->SetBranchAddress("bkg_killer",&bkg_killer);
    
    //loop over the entries
    for(int iEntry; iEntry<tot_entries; ++iEntry){
        
        tree->GetEntry(iEntry);
        
        //skip all candidates below the optimal cut point
        if(bkg_killer<0.92) continue;
        
        //fill the histograms
        h_data->Fill(B_de);
    }
    
    
    //Let's define the PDF for the fit, using TF1
    //https://root.cern.ch/doc/master/classTF1.html
    
    //The total function that describes our observed distribution
    TF1* pdf = new TF1("pdf","gaus(0)+gaus(3)+pol1(6)",min_de,max_de);
    
    //signal gauss, normalisation constant
    pdf->SetParName  (0,  "N_{sig}");
    pdf->SetParameter(0,  100);
    //signal gauss, mean fixed
    pdf->SetParName  (1,  "#mu_{sig}");
    pdf->FixParameter(1,   0.);
    //signal gauss, std dev fixed
    pdf->SetParName  (2,  "#sigma_{sig}");
    pdf->FixParameter(2,0.015);
    //mis-id gauss, normalisation constant
    pdf->SetParName  (3,  "N_{misid}");
    pdf->SetParameter(3,   10);
    //mis-id gauss, mean fixed
    pdf->SetParName  (4,  "#mu_{misid}");
    pdf->FixParameter(4,0.042);
    //mis-id gauss, std dev fixed
    pdf->SetParName  (5,  "#sigma_{misid}");
    pdf->FixParameter(5,0.015);
    //background intercept and slope
    pdf->SetParName  (6,  "p_{0}^{bkg}");
    pdf->SetParName  (7,  "p_{1}^{bkg}");
    //just some style for drawing...
    pdf->SetLineColor(kBlue-3);
    pdf->SetLineWidth(2);
    
    //and now fit, in the range definined by the histogram (option R)
    //option N = not draw (otherwise it draws a canvas with a plot by default)
    cout << "\n First fit, fixing all possible parameters: \n\n";
    h_data->Fit("pdf","RN");
    
    cout << "\n\n Let's try to release the signal std dev \n\n";
    pdf->ReleaseParameter(2); //signal gauss, std dev fixed
    h_data->Fit("pdf","RN");
    
    cout << "\n\n Update the mis-id std dev \n";
    cout << " And release also the mis-id mean \n";
    pdf->FixParameter(5, pdf->GetParameter(2)); //signal gauss, std dev fixed
    pdf->ReleaseParameter(4);
  
    cout << " and do a binned-likelihood fit, instead of a chi2 \n\n";
    //option L = binned likelihood fit
    //Use FitResultPtr to retreive all information about the fit
    //Need to add the option S
    TFitResultPtr fit = h_data->Fit("pdf","LRS");
    //now we can get covariance matrix. We will store in a TMatrixDSym
	TMatrixDSym cov = fit->GetCovarianceMatrix();
	cov.Print();
  
    //draw the result
    gStyle->SetOptStat(0);
    gStyle->SetOptFit(1111);
    TCanvas* c1 = new TCanvas("c1","c1",600,600);
    
    h_data->SetMinimum(0);
    h_data->SetMarkerColor(kBlack);
    h_data->SetMarkerStyle(8);
    h_data->SetMarkerSize(0.8);
    h_data->SetLineColor(kBlack);
    
    h_data->Draw("err");
    
    //just to draw each component separately...
    //the signal
    TF1* pdf_sig = new TF1("pdf_sig","gaus",min_de,max_de);
    pdf_sig->SetParameters(pdf->GetParameter(0),
                           pdf->GetParameter(1),
                           pdf->GetParameter(2));
    pdf_sig->SetLineColor(kRed);
    pdf_sig->SetLineWidth(2);
    pdf_sig->Draw("same");
    
    //the mis-id B->pipi
    TF1* pdf_misid = new TF1("pdf_misid","gaus",min_de,max_de);
    pdf_misid->SetParameters(pdf->GetParameter(3),
                             pdf->GetParameter(4),
                             pdf->GetParameter(5));
    pdf_misid->SetLineColor(kGreen+3);
    pdf_misid->SetLineWidth(2);
    pdf_misid->Draw("same");
    
    //the background
    TF1* pdf_bkg = new TF1("pdf_bkg","pol1",min_de,max_de);
    pdf_bkg->SetParameters(pdf->GetParameter(6),
                           pdf->GetParameter(7));
    pdf_bkg->SetLineColor(kBlue);
    pdf_bkg->SetLineWidth(2);
    pdf_bkg->SetLineStyle(2);
    pdf_bkg->Draw("same");
    
    TLegend* leg = new TLegend(0.15,0.5,0.4,0.85);
    leg->SetBorderSize(0);
    leg->AddEntry(h_data,"Data","PL");
    leg->AddEntry(pdf,"Fit","L");
    leg->AddEntry(pdf_sig,"B^{0} #rightarrow K^{+}#pi^{-}","L");
    leg->AddEntry(pdf_misid,"B^{0} #rightarrow #pi^{+}#pi^{-}","L");
    leg->AddEntry(pdf_bkg,"background","L");
    leg->Draw();
    
    c1->SaveAs("myFit.pdf");
    c1->SaveAs("myFit.C");
  
    //Get now what we wanted to know!
    double binW = h_data->GetXaxis()->GetBinWidth(1);
    cout << "\n\n From this fit model, \n";
    cout << "Candidate in data histogram: " << h_data->Integral() << endl;
    cout << "Total candidates from fit  : " << pdf->Integral(min_de,max_de)/binW  << endl;
    cout << "Signal B->Kpi candidates   : " << pdf_sig->Integral(min_de,max_de)/binW << endl;
    cout << "Mis-id B->pipi candidates  : " << pdf_misid->Integral(min_de,max_de)/binW << endl;
    cout << "Background candidates      : " << pdf_bkg->Integral(min_de,max_de)/binW << endl;

	cout << "======================================================" << endl;
    cout << " Let's calculate the final result with its uncertainty " << endl;
    //We can calculate the signal yield directly from the fit parameters
    //of the signal pdf (gauss function)
	double Nsig  = fit->Parameter(0);
	double sigma = fit->Parameter(2);
	double S = Nsig*sqrt(2.*TMath::Pi())*sigma/binW;
  
	// and propagate the uncertainty
	double errS = S * sqrt(cov(0,0)/Nsig/Nsig + cov(2,2)/sigma/sigma + 2*cov(0,2)/Nsig/sigma);
  
    printf("The measurement of the signal yield is %.0f +- %.0f \n", S, errS);
    //for curiosity let's print the correlation between Nsig and sigma,
    //is it negligible?
    printf("Corr(Nsig,sigma) = %.3f \n\n",cov(0,2)/sqrt(cov(0,0)*cov(2,2)));
  
	//Let's make a scan of the likelihood
	//as a function of N_sig
	TGraph* scan = new TGraph(40);
	fit->Scan(0,scan,120,220);
	TCanvas* c3 = new TCanvas("c3","c3",600,600);
	scan->SetTitle("NLL scan of N_{sig}; N_{sig}; NLL");
	scan->SetLineColor(kRed);
	scan->SetLineWidth(2);
	scan->Draw("APL");
	c3->SaveAs("myScan.pdf");
  
    //Let's have a look: make 2D confidence regions
    //1sigma contour: region enclosing 68.3% probability
    TGraph* cont1sigma = new TGraph(50);
    fit->Contour(0,2,cont1sigma,0.683);
	cont1sigma->SetLineWidth(2);
    cont1sigma->SetLineColor(kBlue+4);
  
    //2sigma contour: region enclosing 95.5% probability
    TGraph* cont2sigma = new TGraph(50);
    fit->Contour(0,2,cont2sigma,0.955);
    cont2sigma->SetLineStyle(2);
    cont2sigma->SetLineWidth(2);
    cont2sigma->SetLineColor(kBlue+2);
  
    //3sigma contour: region enclosing 99.7% probability
    TGraph* cont3sigma = new TGraph(50);
    fit->Contour(0,2,cont3sigma,0.997);
    cont3sigma->SetLineStyle(3);
    cont3sigma->SetLineWidth(2);
    cont3sigma->SetLineColor(kBlue);
  
    //Draw all together, need to use TMultiGraph
    TCanvas* c2 = new TCanvas("c2","c2",600,600);
    TMultiGraph *mg = new TMultiGraph();
    mg->SetTitle("Confidence regions for #sigma_{sig} vs N_{sig}; N_{sig}; #sigma_{sig} [Gev/c^{2}]");
    mg->Add(cont1sigma,"l");
    mg->Add(cont2sigma,"l");
	mg->Add(cont3sigma,"l");
	mg->Draw("a");
  
    TLegend* lc = new TLegend(0.7,0.7,0.89,0.89);
    lc->SetBorderSize(0);
    lc->AddEntry(cont1sigma,"68.3%","l");
    lc->AddEntry(cont2sigma,"95.0%","l");
    lc->AddEntry(cont3sigma,"99.7%","l");
    lc->Draw();
    c2->SaveAs("myContour.pdf");
  
    
  
    return;
}