// ----------------------------------------------------------------------------------- //
/*
  ROOT macro for illustrating a likelihood fit.

  The example uses as default 100 exponentially distributed random times, with the goal
  of finding the best estimate of the lifetime, tau = 1. Four estimates are considered:
   - Mean
   - Chi-square fit (chi2)
   - Binned Likelihood fit (bllh)
   - Unbinned Likelihood fit (llh)

  While the mean can be simply calculated, the three other methods are based on a scan
  of values for tau in the range [0.5, 2.0]. For each value of tau, the chi2, bllh, and
  llh are calculated (in the two likelihood cases, it is actually -2*log(likelihood)
  which is calculated).

  This is not really an exercise, but more an attempt to illustrate three things:
   - How to make a (binned and unbinned) Likelihood function (and fit).
   - The difference and a comparison between a Chi-square and a (binned) Likelihood.
   - The difference and a comparison between a binned and unbinned Likelihood.
   - What goes on behind the scenes in ROOT, when it is asked to fit something.


  Author: Troels C. Petersen (NBI)
  Email:  petersen@nbi.dk
  Date:   26th of September 2011
*/
// ----------------------------------------------------------------------------------- //

// Include from C++:
#include <fstream>
#include <iostream>
#include <stdio>
#include <stdlib>

// Include from ROOT:
#include <TROOT.h>
#include <TH1.h>
#include <TF1.h>
#include <TCanvas.h>
#include <TStyle.h>
#include <TFile.h>
#include <TMath.h>
#include <TRandom3.h>

//---------------------------------------------------------------------------------- 
double sqr(double a) {
  return a*a;
}



// ----------------------------------------------------------------------------------- //
double func_asympara(double *x, double *par) {
// Parameters are as follows:
//   par[0]: Minimum value (i.e. constant)
//   par[1]: Minimum position (i.e. x of minimum)
//   par[2]: Quadratic term on lower side
//   par[3]: Quadratic term on higher side
// ----------------------------------------------------------------------------------- //
  if (x[0] < par[1])
    return par[0] + par[2]*sqr(x[0]-par[1]);
  else
    return par[0] + par[3]*sqr(x[0]-par[1]);
}



// ----------------------------------------------------------------------------------- //
void LikelihoodFit() {
// ----------------------------------------------------------------------------------- //

  // Set the showing of statistics and fitting results (0 means off, 1111 means all on):
  gROOT->Reset();
  gStyle->SetOptStat(1111);
  gStyle->SetOptFit(1111);

  // Set the graphics:
  gStyle->SetStatBorderSize(1);
  gStyle->SetCanvasColor(0);

  // Parameters of the problem:
  const int Ntimes = 100;     // Number of times
  double tau_truth = 1.0;     // We choose (like Gods!) the lifetime.
  vector<double> t;           // Vector of times

  bool verbose = false;       // Should the program print (a lot) or not? (true/false)
  TRandom3 r;

  int Nbin_exp = 50;          // Number of bins in histogram
  TH1F *Hist_Exp = new TH1F("Hist_Exp", "Hist_Exp", Nbin_exp, 0.0, 10.0);


  // ------------------------------------
  // Generate data:
  // ------------------------------------
  for (int i=0; i < Ntimes; i++) {
    t.push_back(-tau_truth*log(r.Uniform()));       // Exponential with lifetime tau.
    // t.push_back(r.Exp(tau_truth));               // Could also use this line!
    Hist_Exp->Fill(t[i]);
    if (verbose) printf("  %5.2f", t[i]);
  }
  if (verbose) printf("\n\n");


  // ------------------------------------
  // Analyse data:
  // ------------------------------------

  // Define the number of tau values to test in Chi2 and LLH:
  const int Ntau = 1500;
  double min_tau = 0.5;
  double max_tau = 2.0;
  double delta_tau = (max_tau-min_tau) / Ntau;

  // Loop over hypothesis for the value of tau and calculate Chi2 and LLH:
  // ---------------------------------------------------------------------
  double tau[Ntau+1], chi2[Ntau+1], bllh[Ntau+1], llh[Ntau+1];
  double chi2_minval = 999999.0;  double chi2_minpos = 0.0;
  double bllh_minval = 999999.0;  double bllh_minpos = 0.0;
  double llh_minval  = 999999.0;  double llh_minpos  = 0.0;

  for (int itau=0; itau < Ntau+1; itau++) {
    double tau_hypo = min_tau + double(itau)*delta_tau;   // Scan in values of tau
    tau[itau] = tau_hypo;

    // Chi2 and binned likelihood (from loop over bins in histogram):
    chi2[itau] = 0.0;
    bllh[itau] = 0.0;
    for (int ibin=1; ibin < Nbin_exp; ibin++) {
      double xlow_bin = Hist_Exp->GetBinCenter(ibin) - 0.1;
      double xhigh_bin = Hist_Exp->GetBinCenter(ibin) + 0.1;
      double nexp = double(Ntimes) * (exp(-xlow_bin/tau_hypo) - exp(-xhigh_bin/tau_hypo));
      double nobs = Hist_Exp->GetBinContent(ibin);

      if (nobs > 0) chi2[itau] += sqr(nobs-nexp) / nobs;
      bllh[itau] += -2.0*log(TMath::Poisson(int(nobs), nexp));
    }

    // Unbinned likelihood (from loop over events):
    llh[itau]  = 0.0;
    for (int iev=0; iev < Ntimes; iev++) {
      llh[itau] += -2.0*log(1.0/tau_hypo*exp(-t[iev]/tau_hypo));
    }

    if (verbose)
      printf(" %3d. tau = %4.2f   chi2 = %6.2f   log(bllh) = %6.2f   log(llh) = %6.2f \n",
	     itau, tau_hypo, chi2[itau], bllh[itau], llh[itau]);

    // Search for minimum values of chi2, bllh, and llh:
    if (chi2[itau] < chi2_minval) {chi2_minval = chi2[itau]; chi2_minpos = tau_hypo;}
    if (bllh[itau] < bllh_minval) {bllh_minval = bllh[itau]; bllh_minpos = tau_hypo;}
    if (llh[itau] < llh_minval) {llh_minval = llh[itau]; llh_minpos = tau_hypo;}
  }

  printf("  Minimum found:   chi2 = %7.4f    bllh = %7.4f    llh = %7.4f \n",
	 chi2_minpos, bllh_minpos, llh_minpos);



  // ------------------------------------
  // Plot and fit results:
  // ------------------------------------

  // ChiSquare:
  TCanvas* c_chi2 = new TCanvas("c_chi2","",80,40,600,450);
  TGraphErrors* graph_chi2 = new TGraphErrors(Ntau+1, tau, chi2);
  TF1 *fit_chi2 = new TF1("fit_chi2", func_asympara, chi2_minpos - 0.20, chi2_minpos + 0.30, 4);
  fit_chi2->SetParameters(chi2_minval, chi2_minpos, 20.0, 20.0);
  fit_chi2->SetLineColor(4);
  graph_chi2->SetMarkerStyle(20);
  graph_chi2->SetMarkerSize(0.4);
  graph_chi2->Fit("fit_chi2","R");
  graph_chi2->Draw("AP");

  double minval = fit_chi2->GetParameter(0);
  double minpos = fit_chi2->GetParameter(1);
  double err_lower = 1.0 / sqrt(fit_chi2->GetParameter(2));
  double err_upper = 1.0 / sqrt(fit_chi2->GetParameter(3));

  TLine* line_min = new TLine(chi2_minpos - 0.30, minval, chi2_minpos + 0.30, minval);
  line_min->Draw();
  TLine* line_min1 = new TLine(chi2_minpos - 0.30, minval+1.0, chi2_minpos + 0.30, minval+1.0);
  line_min1->Draw();
  TLine* line_low = new TLine(minpos-err_lower, minval+1.0, minpos-err_lower, minval-2.0);
  line_low->Draw();
  TLine* line_upp = new TLine(minpos+err_upper, minval+1.0, minpos+err_upper, minval-2.0);
  line_upp->Draw();

  printf("  Chi2 fit gives:    tau = %6.4f + %6.4f - %6.4f \n", minpos, err_lower, err_upper);

  // Go through tau values to find minimum and +-1 sigma:
  // This assumes knowing the minimum value, and Chi2s above Chi2_min+1
  if (((chi2[0] - chi2_minval) > 1.0) && ((chi2[Ntau] - chi2_minval) > 1.0)) {
    bool found_lower = false;
    bool found_upper = false;
    double tau_lower, tau_upper;
    for (int itau=0; itau < Ntau+1; itau++) {
      if ((!found_lower) && ((chi2[itau] - chi2_minval) < 1.0)) {
	tau_lower = tau[itau];
	found_lower = true;
      }
      if ((found_lower) && (!found_upper) && ((chi2[itau] - chi2_minval) > 1.0)) {
	tau_upper = tau[itau];
	found_upper = true;
      }
    }
    printf("  Chi2 scan gives:   tau = %6.4f + %6.4f - %6.4f \n",
	   chi2_minpos, chi2_minpos-tau_lower, tau_upper-chi2_minpos);
  } else {
    printf("Error: Chi2 values do not fulfill requirements for finding minimum and errors! \n");
  }

  c_chi2->Update();
  // c_chi2->SaveAs("FitMinimum_chi2.png");



  // Binned Likelihood:
  TCanvas* c_bllh = new TCanvas("c_bllh","",110,60,600,450);
  TGraphErrors* graph_bllh = new TGraphErrors(Ntau+1, tau, bllh);
  TF1 *fit_bllh = new TF1("fit_bllh", func_asympara, bllh_minpos - 0.20, bllh_minpos + 0.20, 4);
  fit_bllh->SetParameters(bllh_minval, bllh_minpos, 20.0, 20.0);
  fit_bllh->SetLineColor(4);
  graph_bllh->SetMarkerStyle(20);
  graph_bllh->SetMarkerSize(0.4);
  graph_bllh->Fit("fit_bllh","R");
  graph_bllh->Draw("AP");

  c_bllh->Update();
  // c_bllh->SaveAs("FitMinimum_bllh.png");



  // Unbinned Likelihood:
  TCanvas* c_llh = new TCanvas("c_llh","",140,80,600,450);
  TGraphErrors* graph_llh = new TGraphErrors(Ntau+1, tau, llh);
  TF1 *fit_llh = new TF1("fit_llh", func_asympara, llh_minpos - 0.20, llh_minpos + 0.20, 4);
  fit_llh->SetParameters(llh_minval, llh_minpos, 20.0, 20.0);
  fit_llh->SetLineColor(4);
  graph_llh->SetMarkerStyle(20);
  graph_llh->SetMarkerSize(0.4);
  graph_llh->Fit("fit_llh","R");
  graph_llh->Draw("AP");

  c_llh->Update();
  // c_llh->SaveAs("FitMinimum_llh.png");







  // Fit the data using ROOT (both chi2 and binned likelihood fits):
  c1 = new TCanvas("c1","",180,110,600,450);

  Hist_Exp->SetXTitle("Lifetime [s]");
  Hist_Exp->SetYTitle("Frequency [ev/0.1s]");

  TF1 *rootfit_chi2 = new TF1("rootfit_chi2", "100.0 * 0.20 / [0] * exp(-x/[0])", 0.0, 10.0);
  rootfit_chi2->SetParameter(0, 1.0);
  rootfit_chi2->SetLineColor(4);
  Hist_Exp->Fit("rootfit_chi2","RE");

  TF1 *rootfit_bllh = new TF1("rootfit_bllh", "100.0 * 0.20 / [0] * exp(-x/[0])", 0.0, 10.0);
  rootfit_bllh->SetParameter(0, 1.0);
  rootfit_bllh->SetLineColor(2);
  Hist_Exp->Fit("rootfit_bllh","RLE+");

  c1->Update();
  // c1->SaveAs("GraphFitted.png");


}


// ----------------------------------------------------------------------------------- //
/*

Questions:
----------
1) 


Advanced questions:
-------------------
1) 


*/
//----------------------------------------------------------------------------------