//////////////////////////////////////////////////////////////////////////////////////////////
//These are the (annotated) solutions for the ROOT exercise you were given on Monday//
//See the previous session for the .root files and description of exercise
//Big thank you to the HASCO for the providing the foundations for this problem!
//////////////////////////////////////////////////////////////////////////////////////////////

#include <TTree.h>
#include <TH1F.h>
#include <TLatex.h>
#include <TF1.h>
#include <TFile.h>
#include <TH1D.h>
#include <string>
#include <iostream>
#include <fstream>
#include "TRandom3.h"
#include <TCanvas.h>
#include "TMath.h"
#include <TLegend.h>
#include <TStyle.h>

using namespace std;

/////////////////////////////////////////////////////////////////////////////////////////
//Problem: Search for H-> gamma gamma
//  - We want to perform a search for H -> gamma gamma using the invariant diphoton mass
/////////////////////////////////////////////////////////////////////////////////////////

//Note: Sections work in standalone form
//i.e. if you comment out the "Part e" section, the "Parts a-d" code will still properly run

void Higgs(){
    
// ----------------Part a ----------------- //

// Plot the Data:
//  - Inside PseudoData_Histogram_100fb.root is a TH1D histogram called "signal"
//    (shows the invariant diphoton mass). Plot this! 
    
// --------------------------------------- //  
    
//Read in the Data:    
TFile *histoFile = new TFile("PseudoData_Histogram_100fb.root", "READ");
TH1D  *hData     = (TH1D*) histoFile -> Get("signal");
    
//Display # of Data Events: 
std::cout << hData -> Integral() << std::endl; 
    
//Define a canvas. We will divide this into two pads, since part (c) asks
//for a [Data-MC] difference plot
    
TCanvas *c = new TCanvas("c", "c", 600, 600);
TPad *pad1 = new TPad("pad1","main",  0, 0.3, 1, 1.0); //draw this pad from 30% to 100% of the y canvas
pad1->SetFillColor(0);
pad1->Draw();

TPad *pad2 = new TPad("pad2","ratio",  0, 0.05, 1, 0.3); //draw this pad from 0.05% to 30% of the y canvas                                                
pad2->SetFillColor(0);
pad2->Draw();
pad1->cd();
    
//Plot our data:
hData -> SetLineColor(kBlack);
hData -> SetMarkerColor(kBlack);
hData -> SetMarkerStyle(2);
hData -> SetMarkerSize(1.5);
hData->Draw();
c->Draw();

// ----------------Part b ----------------- //

// Plot the background simulation:
//  - Create MC sig+BG histograms and fill them with invariantMass weighted by the eventWeight
//  - Scale the simulation to the correct integrated luminosity of the data and 
//    compare the data to the background-only hypothesis 
    
// --------------------------------------- //

//Read in the trees in the signal and BG file:
TFile *fSig = new TFile("Signal_1fb.root",     "READ");
TFile *fBkg = new TFile("Background_1fb.root", "READ");
    
TTree *tSig = (TTree*) fSig -> Get("tree");
TTree *tBkg = (TTree*) fBkg -> Get("tree");

//Inside the tree, are two variables "invariantMass" and "eventWeight"
double mass_sig; double eventWeight_sig;
double mass_bkg; double eventWeight_bkg;

//We need to link these branch variables to local variables
tSig -> SetBranchAddress("invariantMass", &mass_sig); //note the &. This means the "address" and is needed here
tSig -> SetBranchAddress("eventWeight",   &eventWeight_sig);

tBkg -> SetBranchAddress("invariantMass", &mass_bkg);
tBkg -> SetBranchAddress("eventWeight",   &eventWeight_bkg);

// define two histograms
TH1D *hSig = new TH1D("signal", "", 50, 100, 160);
TH1D *hBkg = new TH1D("bkg",    "", 50, 100, 160);

int nEntries_Sig = tSig -> GetEntries(); //Need the total # entries for our event loop
int nEntries_Bkg = tBkg -> GetEntries();

//Fill the 1D histograms with the invarientMass weighed by the "eventWeight":
for(int i = 0; i < nEntries_Sig; ++i){
    tSig -> GetEntry(i);
    hSig -> Fill(mass_sig, eventWeight_sig);
}
    
for(int i = 0; i < nEntries_Bkg; ++i){
    tBkg -> GetEntry(i);
    hBkg -> Fill(mass_bkg, eventWeight_bkg);
}

//Output the number of events:
std::cout << "Before reweighting: " << std::endl;
std::cout << "Signal:\t\t" << hSig -> Integral() << "\n" << "Background:\t" << hBkg -> Integral() << std::endl;

//However, as we stated in the question, we need to scale the simulation by x100, since 
//Our data was in 100 fb^-1, while the simulation is in units of 100 fb^-1
hSig -> Scale(100.0);
hBkg -> Scale(100.0);
    
//Output the number of events (after reweighting):
std::cout << "After reweighting: " << std::endl;
std::cout << "Signal:\t\t" << hSig -> Integral() << "\n" << "Background:\t" << hBkg -> Integral() << std::endl;

//Plot the reweighted BG model:
hBkg -> SetLineColor(kBlue);
hBkg -> Draw("HSAME");
c ->Draw();


// ----------------Part c ----------------- //

// Background-only hypothesis test:
//  - Perform a fit to the background in order to get a stable background model 
//  - Compare the data with the background model by plotting the difference 
//    in a sub-plot below the main plot.
    
// --------------------------------------- //

//We want to get a model for our background. To do this, fit it with a function.
// Since it looks very exponential (recall the outputs of (b)), try fitting w/ an exponential
    
TF1 *fit = new TF1("f1", "[0] + exp([2]*x+[1])", -1, 12);
fit -> SetLineColor(kRed);
hBkg -> Fit(fit);
c->Draw(); 
    
//Now we want to find  the difference between data and background model. This will elp us more clearly identify a possible signal.
    
//Let's create a new histogram in which we store the difference:
TH1D *hDiff = (TH1D*)hData -> Clone(0);
int nBins = hData->GetNbinsX();

//Loop over all bins, and find the [Data-fit model] difference
// Get Function info by using : fit -> Eval(binCenter);
// Get Histogram info by using: hData->GetXaxis()->GetBinCenter(iBin); hData->GetBinContent(iBin);
    
for(int iBin = 1; iBin <= nBins; ++iBin){
    double binCenter   = hData->GetXaxis()->GetBinCenter(iBin);
    double dataValue   = hData->GetBinContent(iBin);
    double functionVal = fit -> Eval(binCenter);
    double difference  = dataValue - functionVal;
    hDiff -> SetBinContent(iBin, difference); 
}
//Plot this in our second (lower) pad
pad2 -> cd();
hDiff -> Draw();
//Draw new canvas with both pads
c -> Draw();
    
//To more clearly show fluctuations, I've included a dotted line through 0:
TF1 *zero = new TF1("zero", "0", 100, 160);
zero -> SetLineColor(kBlack);
zero -> SetLineWidth(3);
zero -> SetLineStyle(2);
zero -> Draw("SAME");
c->Draw();

// ----------------Part d ----------------- //

// Finalize the plot:
//   - Add the signal simulation to the background histogram.
//   - Make the plot look nice.
    
// --------------------------------------- //

//Add the signal to our background model:        
hBkg -> Add(hSig);
c -> Draw();
    
//Now add more features to the plot (make it look nicer):    

//Add axis labels:
gStyle->SetOptStat(0);
hData -> GetYaxis() -> SetTitle("Number of events");
hData -> GetYaxis() -> SetTitleOffset(1.55);
hData -> GetYaxis() -> SetTitleSize(20);
hData -> GetYaxis() -> SetTitleFont(43);
hData -> GetYaxis() -> SetLabelSize(15);
hData -> GetYaxis() -> SetLabelFont(43);

hDiff -> GetYaxis() -> SetTitle("Data-Fit");
hDiff -> GetYaxis() -> SetTitleSize(20);
hDiff -> GetYaxis() -> SetTitleFont(43);
hDiff -> GetYaxis() -> SetTitleOffset(1.55);
hDiff -> GetYaxis() -> SetLabelSize(15);
hDiff -> GetYaxis() -> SetLabelFont(43);
hDiff -> GetYaxis() ->SetRangeUser(-500, 700);

hDiff -> GetXaxis() -> SetTitle("m_{#gamma #gamma} [GeV]");
hDiff -> GetXaxis() -> SetTitleSize(20);
hDiff -> GetXaxis() -> SetTitleFont(43);
hDiff -> GetXaxis() -> SetTitleOffset(4.);
hDiff -> GetXaxis() -> SetLabelSize(15);
hDiff -> GetXaxis() -> SetLabelFont(43);

c->Draw();

// Connect the two pads together (so they have no gap in between):
pad1 -> SetBottomMargin(0);
pad2 -> SetTopMargin(0);
pad2 -> SetBottomMargin(0.2);
c -> Draw();

//Add a legend:
pad1->cd();
TLegend *fLegend = new TLegend(0.625, 0.66, 0.625+0.255, 0.865);
fLegend -> AddEntry(hData,   "PseudoData", "lp");
fLegend -> AddEntry(hBkg,    "MC",         "l");
fLegend -> AddEntry(fit,     "Fit (BG Only) ",        "l");
fLegend -> SetFillColor(0);
fLegend -> SetLineColor(0);
fLegend -> SetBorderSize(0);
fLegend -> SetTextFont(72);
fLegend -> SetTextSize(0.035);
fLegend -> Draw("SAME");
c->Draw();

//Add some text to the figure itself:
    
TLatex l1;
l1.SetTextAlign(9);
l1.SetTextSize(0.04);
l1.SetNDC();
l1.DrawLatex(0.21, 0.160, "#int L dt = 100 fb^{-1}");

TLatex l2;
l2.SetTextAlign(9);
l2.SetTextFont(72);
l2.SetTextSize(0.04);
l2.SetNDC();
l2.DrawLatex(0.21, 0.310, "EIEIOO");

TLatex l3;
l3.SetTextAlign(9);
l3.SetTextSize(0.04);
l3.SetNDC();
l3.DrawLatex(0.21, 0.235, "Simulation");
c->Draw();


// ----------------Part e ----------------- //

// Perform Signal strength extraction:
//  - Perform a signal + background fit. 
//  - From the signal component of this fit, determine the number of signal events,
//    and compare to the number we expect from the signal simulation. 
    
// --------------------------------------- //

//Peform a signal + background fit The signal looks Gaussian, so let's try using that:
TF1 *fit_sig_bkg = new TF1("fit_sig_bkg", "[0] + exp([2]*x+[1]) + [3]*TMath::Gaus(x,[4],[5])", 100, 160);
    
//In order to help the fit converge, now that it has quite a lot of parameters, 
// we can initialise the parameters to values that are about right
    
// (Note: you should always try to give reasonable initial estimates to your parameters,
//  otherwise your fit may have stability issues and may not converge)
    
fit_sig_bkg->SetParameter(0, fit->GetParameter(0));
fit_sig_bkg->SetParameter(1, fit->GetParameter(1));
fit_sig_bkg->SetParameter(2, fit->GetParameter(2));
fit_sig_bkg->SetParameter(4, 125);
fit_sig_bkg->SetParameter(5, 3);
    
//Peform the fit and draw it on the plot
pad1 -> cd();
fit_sig_bkg -> SetLineColor(kBlack);
hData -> Fit(fit_sig_bkg);
fLegend -> AddEntry(fit_sig_bkg, "Fit (Signal+ BG)", "l");
fLegend -> Draw(); // redraw the legend
    
//Output the fit Parameters ( this is done using the GetParameter() function )
cout << "Fitted signal normalisation: " << fit_sig_bkg->GetParameter(3) << " +- " << fit_sig_bkg->GetParError(3) << endl;
cout << "Fitted signal mean: "          << fit_sig_bkg->GetParameter(4) << " +- " << fit_sig_bkg->GetParError(4) << endl;
cout << "Fitted signal width: "         << fit_sig_bkg->GetParameter(5) << " +- " << fit_sig_bkg->GetParError(5) << endl;

//The integral over all x of a gaussian is = sqrt(pi). 
//So For a Gaussian of width s and normalisation n, it is therefore just s * n * sqrt(pi)
    
//And we can calculate the number of signal events:    
float fittedSignalEvents = fit_sig_bkg->GetParameter(3) * fit_sig_bkg->GetParameter(5) * TMath::Sqrt(TMath::Pi());

//Compare our number of signal events to what we expect from our simulation. 
//This is what's called out signal strength (mu)
//mu = number of events we extract above our background prediction 
//     divided by the number of signal events predicted by our model. 
//It should be close to 1 if our prediction is reasonable
    
float mu = fittedSignalEvents / hSig -> Integral();
cout << "Signal strength: " << mu << endl;

//Now plot the fitted signal in the ratio plot (so we can see it by eye):
pad2 -> cd();
TF1 * fit_sig = new TF1("fit_sig","gaus(0)",100, 160);
fit_sig->SetParameter(0, fit_sig_bkg->GetParameter(3));
fit_sig->SetParameter(1, fit_sig_bkg->GetParameter(4));
fit_sig->SetParameter(2, fit_sig_bkg->GetParameter(5));
fit_sig->SetLineColor(kBlack);
fit_sig->Draw("same");
    
//Congratulations! You've just found the Higg's!
//Now that you've done this, you are ready to tackle your own analyses this summer! Good luck!

}