#include "TBox.h"
#include "TArrow.h"
#include "TCanvas.h"
#include "TH1D.h"
#include "TH2D.h"
#include "TLine.h"
#include "TMath.h"
#include "TCanvas.h"
#include "TRandom3.h"
#include "TROOT.h"
#include "TLatex.h"
#include "TFile.h"
#include "TLegend.h"
#include "TLegendEntry.h"
#include "Math/QuantFuncMathCore.h"
#include "Math/ProbFuncMathCore.h" // for ROOT::Math::gaussian_cdf

#include <iostream>
using namespace std;

//-------------------------------------------------------------------------------------------------------------------------
//-- full functions
TH1D * GetMassDistribution(int Itype = 1, double scalefactor = 1.00);
void MassPlot(int Irebin = 20);

//-- skeleton functions
void SideBandFit(int Irebin = 10);
double Get_TestStatistic(TH1D *h_mass_dataset, TH1D *h_template_bgr, TH1D *h_template_sig);
TH1D * GenerateToyDataSet(TH1D *h_mass_template);

//-- some fusefull functions
double IntegratePoissonFromRight(double mu, int N_obs);
double IntegrateFromRight(TH1D * h_X_bgr, double X_value);
vector<double> Get_Quantiles( TH1D* hist );
void AddText( Double_t txt_x = 0.50, Double_t txt_y = 0.50, const char * txt = "dummy", Double_t txt_size = 0.045, 
	      Double_t txt_angle = 0., const char * Alignment = "left", Int_t UseNormalizedSize = 1, Int_t txt_color =1 );
//-------------------------------------------------------------------------------------------------------------------------



//========================================
// S O M E   F I N A L   F U N C T I O N S 
//========================================


//========================================================
TH1D * GetMassDistribution(int Itype, double scalefactor){
//========================================================
 //----------------------------------------------------------
 // Goal: return the histogram of the 4-lepton invariant mass
 //       for given type with an optional scale factor
 //
 //       Itype 1 = ZZ SM background
 //             2 = data
 //           125 = Higgs 125
 //           200 = Higgs 200
 //
 //      scalefactor: histograms will be scaled with this number
 //
 //  Note: Histograms have ~200 MeV bins, so need to rebin
 //---------------------------------------------------------

  //-- [1] Get histogram from the file
  TH1D *h_mass = 0;
  TDirectory* dir = gDirectory;   
  TFile *file = new TFile("Histograms_fake.root", "READ");
  dir->cd();

  //-- Higgs 125
  if(Itype == 125){
    h_mass  = (TH1D*) file->Get("h_m4l_Higgs125_fake")->Clone("h_mass");     
  }
  //-- Higgs 200
  if(Itype == 200){
    h_mass  = (TH1D*) file->Get("h_m4l_Higgs200_fake")->Clone("h_mass");     
  }
  //-- ZZ SM background
  if(Itype == 1){
    h_mass  = (TH1D*) file->Get("h_m4l_ZZ_fake")->Clone("h_mass");     
  }
  //-- data
  if(Itype == 2){
    h_mass  = (TH1D*) file->Get("h_m4l_data_fake")->Clone("h_mass");     
  }
  
  //-- [2] scale histograms
  int Nbins = h_mass->GetNbinsX();
  for (int i_bin = 1; i_bin < Nbins; i_bin++){
    double mu_bin = h_mass->GetBinContent(i_bin);
    h_mass -> SetBinContent( i_bin, scalefactor * mu_bin);
  }

  //-- [3] return histogram   
  return h_mass;

  //===========================
} // end GetMassDistribution()
  //===========================




//========================
void MassPlot(int Irebin){
//========================
  // ------------------------------------------
  // Goal: produce SM+Higgs+data plot
  //       Note: rebinning is only for plotting
  // ------------------------------------------

  //------------------------------------
  //-- Standard stuff and prepare canvas
  //------------------------------------
  gROOT->Clear();
  gROOT->Delete();

  //-- Prepare canvas and plot histograms
  TCanvas * canvas1 = new TCanvas( "canvas1","Standard Canvas",600,400);  
  canvas1->SetLeftMargin(0.125);
  canvas1->SetBottomMargin(0.125); 
  canvas1->cd(); 

  //------------------------------------------------------------------
  //-- [1] Prepare histograms
  //--     o Get histograms from the files (signal, background and data)
  //--     o Make cumulative histograms (for signal and background)
  //------------------------------------------------------------------

  //-- Get histograms from the files (higgs, zz and data)
  TH1D *h_sig, *h_bgr, *h_data;
  h_sig  = GetMassDistribution(125);
  h_bgr  = GetMassDistribution(1);
  h_data = GetMassDistribution(2);

  //-----------------------------------
  //-- [2] Plot histograms and make gif
  //--     o rebin histograms 
  //--     o prepare cumulative histogram
  //--     o make plot + opsmuk + gif
  //-----------------------------------

  //-- Rebin histograms (only for plotting)
  h_sig->Rebin(Irebin);
  h_bgr->Rebin(Irebin);
  h_data->Rebin(Irebin);

  //-- Prepare cumulative histogram for signal + background 
  TH1D *h_sig_plus_bgr = (TH1D* ) h_bgr->Clone("h_sig_plus_bgr");
  h_sig_plus_bgr->Reset();
  for (int i_bin = 1; i_bin < h_bgr->GetNbinsX(); i_bin++){
       h_sig_plus_bgr->SetBinContent( i_bin, h_sig->GetBinContent(i_bin) + h_bgr->GetBinContent(i_bin));
       printf("  REBINNED HISTOGRAM:  bin %d, Ndata = %d\n",i_bin,(int)h_data->GetBinContent(i_bin));
  }

  //-- prepare histograms and plot them on canvas
  double Data_max = h_data->GetBinContent(h_data->GetMaximumBin());
  double Ymax_plot = 1.10* (Data_max + TMath::Sqrt(Data_max));
  h_sig_plus_bgr->SetFillColor(7); 
  h_sig_plus_bgr->SetAxisRange(0.,Ymax_plot,"Y");
  h_sig_plus_bgr->SetAxisRange(0.,400.,"X");
  h_bgr->SetFillColor(2); 
  h_sig_plus_bgr->Draw("hist");  
  h_bgr->Draw("same");  
  h_bgr->Draw("axis same");  
  h_data->Draw("e same");

  //-- some nice axes and add legend
  AddText( 0.900, 0.035, "4-lepton invariant mass [GeV]",0.060, 0.,"right");                             // X-axis
  AddText( 0.040, 0.900, Form("Number of events / %3.1f GeV",h_bgr->GetBinWidth(1)) ,0.060,90.,"right"); // Y-axis
  TLegend *leg1 = new TLegend(0.65,0.65,0.90,0.85);
  leg1->SetBorderSize(0); leg1->SetFillColor(0);
  TLegendEntry *leg1a = leg1->AddEntry(h_bgr,          " SM(ZZ)", "f");  leg1a->SetTextSize(0.04);
  TLegendEntry *leg1b = leg1->AddEntry(h_sig_plus_bgr, " Higgs" , "f");  leg1b->SetTextSize(0.04);
  leg1->Draw();

  //-- prepare gif
  canvas1->Print(Form("./MassPlot_rebin%d.gif",Irebin));

  return;

   //===============
 } // end MassPlot()
   //===============




//===============================================
// S O M E   S K E L E T O N    F U N C T I O N S 
//===============================================


//=============================================================
void Significance_Optimization(double Lumi_scalefactor = 1.00){
//=============================================================

  printf("\n Significance_Optimization()\n\n");

  //------------------------------------------------------------------
  //-- [1] Prepare histograms
  //--     o Get histograms from the files (signal, background and data)
  //--     o scale to correct luminosity
  //------------------------------------------------------------------

  //-- Get histograms from the files (higgs, zz and data)
  TH1D *h_sig, *h_bgr, *h_data;
  printf ("\n  INFO: Mass distribution in the 4 lepton channel\n");
  h_sig  = GetMassDistribution(125, Lumi_scalefactor);
  h_bgr  = GetMassDistribution(1, Lumi_scalefactor);
  h_data = GetMassDistribution(2, Lumi_scalefactor);

  //-------------------------------------------------------
  //-- [2] Compute significance for various mass windows
  //--     o try various options for window (use histogram)
  //--     o compute expected and observed significance 
  //-------------------------------------------------------

  //-- Define histogram that defines mass windows to try
  TH1D *h_masswindow          = new TH1D("h_masswindow","",250,0.,25.);           // make a mass window - full width between 0 and 25 GeV
  TH1D *h_masswindow_expected = new TH1D("h_masswindow_expected","",250,0.,25.);  // histogram to hold results for expected
  TH1D *h_masswindow_observed = new TH1D("h_masswindow_observed","",250,0.,25.);  // histogram to hold results for observed

  //---------------------------------
  //-- Loop over various mass windows
  //---------------------------------
  for (int i_bin = 1; i_bin<h_masswindow->GetNbinsX(); i_bin++ ){
  
    //-- get full width of mass window (bin center of our histogram) and the number of events in mass window for each event type
    double masswindow_fullwidth = h_masswindow->GetBinCenter(i_bin);    
    printf("   Trying as mass window: %5.2f GeV\n", masswindow_fullwidth);

    //-- [a] determine the number of events in the mass window for each event type
    //       Ndata_win, Nbgr_win and Nsig_win

    //-- [b] compute EXPECTED significance and save in histogram
    double pvalue_expected       = 0.75; // you need to do this yourself
    double significance_expected = ROOT::Math::gaussian_quantile_c(pvalue_expected,1);
    h_masswindow_expected->SetBinContent(i_bin, significance_expected);

    //-- [c] compute OBSERVED significance and save in histogram
    double pvalue_observed       = 0.25; // you need to do this yourself
    double significance_observed = ROOT::Math::gaussian_quantile_c(pvalue_observed,1);
    h_masswindow_observed->SetBinContent(i_bin, significance_observed);

    printf("           expected significance = 1.000 || observed significance = 1.000 \n");

  } // end loop over width mass window

  //-- print optimum to the screen
  //


  //----------------------------------
  //-- [3] Plot histogram and make gif
  //----------------------------------
  TCanvas * canvas1 = new TCanvas( "canvas1","Standard Canvas",600,400);  
  canvas1->SetLeftMargin(0.125);
  canvas1->SetBottomMargin(0.125); 
  canvas1->cd(); 
 
  h_masswindow_expected->SetLineColor(1);
  h_masswindow_expected->SetLineWidth(2);
  h_masswindow_observed->SetLineColor(4);
  h_masswindow_observed->SetLineWidth(2);

  h_masswindow_expected->SetAxisRange(-1.,6.,"Y");
  h_masswindow_expected->Draw("l");
  if(fabs(Lumi_scalefactor-1.00)<0.01){
    h_masswindow_observed->Draw("l same");
  }
  //-- axes
  AddText( 0.900, 0.035, "Mass window GeV",0.060, 0.,"right"); // X-axis
  AddText( 0.040, 0.900, "Significance" ,0.060,90.,"right");   // Y-axis    
  AddText( 0.225, 0.825, Form("Luminosity scalefactor = %5.1f",Lumi_scalefactor),0.050, 0.,"left");            

  AddText( 0.700, 0.200, "Expected significance",0.050, 0.,"right",1,1);                        
  if(fabs(Lumi_scalefactor-1.00)<0.01){
    AddText( 0.700, 0.300, "Observed significance",0.050, 0.,"right",1,4);                        
  }
  //-- prepare gif
  canvas1->Print(Form("./Significance_Optimization_lumiscalefactor%d.gif",int(Lumi_scalefactor)));

  return;

  //================================
} // end Significance_Optimization()
  //================================






//===========================
void SideBandFit(int Irebin){
//===========================

  printf("\n SideBandFit()\n\n");

  //-------------------------
  //-- [1] Prepare histograms
  //-------------------------
  TH1D *h_bgr, *h_data;
  h_bgr  = GetMassDistribution(1);
  h_data = GetMassDistribution(2);
 
  //-- rebin histograms if necessary
  h_bgr->Rebin(Irebin);
  h_data->Rebin(Irebin);
  printf(" INFO: Rebinning the histograms with a factor %d. Binwidth is now %5.2f GeV\n", Irebin, h_data->GetBinWidth(1));

  //-----------------------------------------
  //-- [2] Loop over scale factor (alpha_bgr)
  //-----------------------------------------
  TH1D *h_scalefactor_bgr = new TH1D("h_scalefactor_bgr","",100,0.1,3.1); // you'll want to put more here I guess
  double scalefactor_bgr  = 1.00; 

  //---------------------------------------------
  // [2a] Loop 1: loop over scalefactors in alpha
  //---------------------------------------------
  for (int i_bin_sf = 1; i_bin_sf <=h_scalefactor_bgr->GetNbinsX(); i_bin_sf++){

    //-- determine the scale factor for the background
    scalefactor_bgr = h_scalefactor_bgr->GetBinCenter(i_bin_sf);
    printf(" Loop 1: I am now trying alpha = %5.2f\n",scalefactor_bgr);
  
    //-----------------------------------------------------------------------------------
    // [2b] Loop 2: loop over bins in the histogram, compute loglik and save in histogram
    //-----------------------------------------------------------------------------------
    double loglik = 1e-10;
    for (int i_bin = 1; i_bin <=  h_data->GetNbinsX(); i_bin++){
      printf("        bin %d, m = %5.2f, Ndata = %5.2f, Nbgr = %5.2f\n",i_bin, h_bgr->GetBinCenter(i_bin),h_bgr->GetBinContent(i_bin),h_data->GetBinContent(i_bin));
      loglik += 1.; // you'll have to do this yourself;  	
    } // end loop over bins

    h_scalefactor_bgr->SetBinContent(i_bin_sf,-2.*loglik);   
  } // end loop over scale factors for the background


    //----------------------------------------------------
    //-- [3] Interpret the -2Log (Likelihood distribution)
    //----------------------------------------------------
    TH1D *h_scalefactor_bgr_rescaled  = (TH1D*) h_scalefactor_bgr->Clone("h_scalefactor_bgr_rescaled");  
    h_scalefactor_bgr_rescaled->Reset();   
    
    //-- Find minimum
    // todo

    //--Rescale and find +/- 1 sigma errors
    // todo

    //-- print summary to screen
    // todo

  return;

  //==================
} // end SideBandFit()
  //==================



//=========================================================================================
double Get_TestStatistic(TH1D *h_mass_dataset, TH1D *h_template_bgr, TH1D *h_template_sig){
//=========================================================================================
  printf(" dummy = %d\n",h_mass_dataset->GetNbinsX() + h_template_bgr->GetNbinsX() + h_template_sig->GetNbinsX());
   
  double test_statistic = 0.;

  //-- do likelihood fit
  //for (int i_bin = 1; i_bin <= Nbins; i_bin++){
      // loglik_bgr +=  // likelihood for the mu=0 (b-only) scenario
      // loglik_sig_plus_bgr +=  // likelihood for the mu=1 (s+b) scenario
  //} // end loop over bins
 
  //-- compute test statistic
  //double test_statistic  =  // find this out yourself

  //-- return test_statistic
  return test_statistic;

} // end Get_TestStatistic()





//===============================================
TH1D * GenerateToyDataSet(TH1D *h_mass_template){
//===============================================
  //-------------------------------------------------
  // Goal: Generate Toy data set from input histogram
  // How:  dumb way -> draw random Poisson in each bin       
  //--------------------------------------------------
  //TRandom3 *R = new TRandom3(0);

  //-- Create new histogram for the data-set
  TH1D *h_mass_toydataset = (TH1D*) h_mass_template->Clone("h_mass_toydataset"); h_mass_toydataset->Reset();
 
  //-- Loop over bins and draw Poisson number of event in each bin
    
  //-- return histogram of toy data-set
  return h_mass_toydataset;
  
} // end GenerateToyDataSet()

















//================================================
// S O M E   U S E F U L L   F U N C T I O N S 
//================================================



//====================================================
double IntegratePoissonFromRight(double mu, int N_obs){
//====================================================
// --------------------------------------------------------------------
// Compute p-value for case zero background uncertainty, i.e. 
//         just integrate Poisson from the right from N_obs to infinity
// --------------------------------------------------------------------

  double integral = 1.; 
  for(int i_obs = 0; i_obs < N_obs; i_obs++){
    integral -= TMath::Poisson(i_obs,mu);
  }
  
  return integral;
  
} // end IntegratePoissonFromRight()


//========================================================
double IntegrateFromRight(TH1D * h_X_bgr, double X_value){
//========================================================
// --------------------------------------------------------------------
// Compute p-value: integrate number of events from X_value to infinity
// --------------------------------------------------------------------

  //-- Integrate distributions
  int Nbins = h_X_bgr->GetNbinsX();
  int X_bin = h_X_bgr->FindBin(X_value); 

  //-- Compute integral from X-value to infinity
  double pvalue = h_X_bgr->Integral(X_bin,Nbins) / h_X_bgr->Integral();
  
  return pvalue;

} // end IntegrateFrom Right()




//=========================================
vector<double> Get_Quantiles( TH1D* hist ){
//=========================================
// Quantiles returns a vector<double> with 5 entries.
// Entries 0 and 4 are the values on the histogram x-axis
// so that 95% of the content lies between these values.
// Entries 1 and 3 bracket 68% in the same way.
// Entry 2 is the median of the histogram.

  //-- define quantiles
  double fraction_1sigma = ROOT::Math::gaussian_cdf(-1.,1.,0.); // 15.8655 %
  double fraction_2sigma = ROOT::Math::gaussian_cdf(-2.,1.,0.); //  2.2750 %
  double probs[5] = {fraction_2sigma, fraction_1sigma, 0.50, 1.00-fraction_1sigma, 1.00-fraction_2sigma };

  //-- output of the quantiles
  double Xvalues[5];

  //-- extract quantiles
  hist->GetQuantiles( 5, Xvalues, probs );
  
  vector<double> Xvalues_output(5);
  for (int i=0; i<5; i++) 
    {
      Xvalues_output[i] = Xvalues[i];
    }

  return Xvalues_output;
} // end Get_Quantiles()





//=======================================================================================================================
void AddText( Double_t txt_x, Double_t txt_y, const char * txt, Double_t txt_size,
              Double_t txt_angle, const char * Alignment, Int_t UseNormalizedSize, Int_t txt_color)
//=======================================================================================================================
{
  Int_t txt_align = 12;
  if ( !strcmp(Alignment, "left"))   { txt_align = 12; } // left 
  if ( !strcmp(Alignment, "right"))  { txt_align = 32; } // right
  if ( !strcmp(Alignment, "center")) { txt_align = 22; } // center

  TLatex* t1 = new TLatex( txt_x, txt_y, txt);
  if(UseNormalizedSize) {t1->SetNDC(kTRUE);} // <- use NDC coordinate
  t1->SetTextSize(txt_size);
  t1->SetTextAlign(txt_align);
  t1->SetTextAngle(txt_angle);
  t1->SetTextColor(txt_color);
  t1->Draw();

} // end AddText()

