// File:  inv14.cc
// Glen Cowan
// RHUL Physics
// Statistics problem for Invisibles 2014

#include <iostream>
#include <string>
#include <vector>
#include <cmath>
#include <TRandom3.h>
#include <TH1D.h>
#include <TFile.h>
#include <TMath.h>

using namespace std;

double lnL(int n, double bTilde, double s, double b, double sigma_bTilde){
  double nd = static_cast<double>(n);
  double u = (b - bTilde)/sigma_bTilde;
  double f = nd*log(s+b) - (s+b) - 0.5*u*u;
  return f;
}

double q0(int n, double bTilde, double sigma_bTilde) {
  double nd = static_cast<double>(n); 
  double sHat = nd - bTilde;
  double bHat = bTilde;
  double sigb2 = sigma_bTilde*sigma_bTilde;
  double bHatHat0 = 0.5*(bTilde - sigb2) + 
     0.5 * sqrt( (bTilde - sigb2)*(bTilde - sigb2) + 4.*sigb2*nd );
  double qZero = 0.;
  if ( sHat > 0. ) {
    double lnLambda0 = lnL(n, bTilde, 0., bHatHat0, sigma_bTilde) 
	             - lnL(n, bTilde, sHat, bHat, sigma_bTilde); 
    qZero = -2. * lnLambda0;
  }
  return qZero;
}

int main() {

  TFile* histFile = new TFile("inv14.root", "recreate");
  TH1D* h_q0 = new TH1D("h_q0", "q0",  80,  0, 40.);

  int seed = 12345;
  TRandom3* ran = new TRandom3(seed);

  double s = 0.;
  double sigma_bTilde = 1.0;
  double b = 6.0;

  // First assume specific values of n and bTilde; compute
  // q0 and use asymptotic formulae to get the significance Z

  int nObs = 12;
  double bTilde = 6.0;
  double q0_obs = q0(nObs, bTilde, sigma_bTilde);
  double Z_obs_wilks = sqrt(q0_obs);
  double p0_obs_wilks = 1. - TMath::Freq(Z_obs_wilks);
  cout << "Observed value of q0      = " << q0_obs << endl;
  cout << "Asymptotic significance Z = " << Z_obs_wilks << endl;
  cout << "Corresponding p-value     = " << p0_obs_wilks << endl;


  // Now generate toy MC experiments to check asymptotic result

  int numWorse = 0;
  int numExp = 10000000;              // number of Toy MC experiments
  for (int i=0; i<numExp; i++){
    int n = ran->Poisson(s+b);
    double bTilde = ran->Gaus(b, sigma_bTilde);
    double q0_gen = q0(n, bTilde, sigma_bTilde);
    h_q0->Fill(q0_gen);
    if ( q0_gen >= q0_obs ) {  numWorse++; }  // count exp with >= compatibility

  }
  double p0 = static_cast<double>(numWorse) / static_cast<double>(numExp);
  double Z0 = TMath::NormQuantile(1. - p0);
  cout << "Significance from MC = " << Z0 << endl;
  cout << "Corresponding p-vaue = " << p0 << endl;

  // Save histogram of q0 (plot with macro plot_q0.C)

  histFile->Write();
  histFile->Close();

  return 0;

}