#include <iostream>
#include <sstream>

#include "TFile.h"
#include "TH1D.h"
#include "TProfile.h"
#include "TCanvas.h"
#include "TString.h"
#include "TStyle.h"
#include "TChain.h"
#include "TH2.h"
#include "TH1.h"
#include "TF1.h"
#include "TTree.h"
#include "TKey.h"
#include "Riostream.h"
#include "TMath.h"
#include "TRandom.h"
#include "TRandom3.h"
#include "TVirtualFitter.h"

#include <stdio.h>
#include <math.h>

// Simple routine needed to calculate the inverse error function
// -------------------------------------------------------------
Double_t ErfInverse (Double_t x);

// Macro that computes p-value and Z-value of N observed vs B predicted
// Poisson counts
// --------------------------------------------------------------------
void Poisson_prob_fluct (double B, double SB, double N, int opt=1) {

  double Niter=10000;

  if (opt!=0 && opt!=1) {
    cout << "Please put fourth argument either =0 (Gaussian nuisance)" << endl;
    cout << "or =1 (LogNormal nuisance)" << endl;
    return;
  }

  int maxN = N*2;
  TH1D * Pois = new TH1D ("Pois", "", maxN, -0.5, maxN-0.5);
  TH1D * PoisGt = new TH1D ("PoisGt", "", maxN, -0.5, maxN-0.5);

  // We throw a random Gaussian smearing SB to B, compute P,
  // and iterate Niter times; we then study the distribution
  // of p-values, extracting the average 

  double Psum=0;
  TH1D * Pdistr = new TH1D ("Pdistr", "", 100, -10., 0.);
  TH1D * TB = new TH1D ("TB", "",100, B-5*SB,B+5*SB);

  double mu, sigma;
  if (opt==0) { // nornal
    mu = B;
    sigma = SB;
  } else { // lognormal
    mu = log(B); // median! omitting the convexity correction -sigma*sigma/2;
    sigma = SB/B;
  }

  for (int iter=0; iter<Niter; iter++) {
    // Extract B from G(B,SB)
    double thisB = gRandom->Gaus(mu,sigma);  // normal
    if (opt==1) thisB = exp(thisB); // lognormal

    TB->Fill(thisB);
    if (thisB<=0) thisB=0.;
    double sum=0.;
    double fact=1.;
    for (int i=0; i<maxN || (opt==0 && i<B+6*SB) || (opt==1 && i<mu+10*sigma); i++) {
      if (i>1) fact*=i;
      double poisson = exp(-thisB)*pow(thisB,i)/fact;
      if (i<N) sum+= poisson;
      Pois->Fill((double)i,poisson);
      if (i>=N) PoisGt->Fill((double)i,poisson);
    }
    double thisP=1-sum;
    if (thisP>0) Pdistr->Fill(log(thisP));
    Psum+=thisP;
  }
  double P = Psum/Niter;
  double Z = sqrt(2) * ErfInverse(1-2*P);

  cout << "Expected P of observing N=" << N << " or more events if B=" 
       << B << "+-" << SB << " : P= " << P << endl;
  cout << "This corresponds to " << Z << " sigma for a Gaussian one-tailed test." << endl;

  Pois->SetLineWidth(3);
  PoisGt->SetFillColor(kBlue);
  TCanvas* T = new TCanvas ("T","Poisson distribution", 500, 500);
  T->Divide(2,2);
  T->cd(1);
  Pois->DrawClone();
  PoisGt->DrawClone("SAME");
  T->cd(2);
  T->GetPad(2)->SetLogy();
  Pois->DrawClone();
  PoisGt->DrawClone("SAME");
  T->cd(3);
  Pdistr->DrawClone();
  T->cd(4);
  TB->Draw();

  delete Pois;
  delete PoisGt;
  delete Pdistr;

}

Double_t ErfInverse(Double_t y) {
// returns the inverse error function obtained
// y must be <-1<y<1

Int_t kMaxit = 50;
Double_t kEps = 1e-14;
Double_t kConst = 0.8862269254527579; // sqrt(pi)/2.0

if(TMath::Abs(y) <= kEps) return kConst*y;

// Newton iterations
Double_t erfi, derfi, Y0,Y1,DY0,DY1;
if(TMath::Abs(y) < 1.0) {
erfi = kConst*TMath::Abs(y);
Y0 = TMath::Erf(0.9*erfi);
derfi = 0.1*erfi;
for (Int_t iter=0; iter<kMaxit; iter++) {
Y1 = 1. - TMath::Erfc(erfi);
DY1 = TMath::Abs(y) - Y1;
if (TMath::Abs(DY1) < kEps) {if (y < 0) return -erfi; else return erfi;}
DY0 = Y1 - Y0;
derfi *= DY1/DY0;
Y0 = Y1;
erfi += derfi;
if(TMath::Abs(derfi/erfi) < kEps) {if (y < 0) return -erfi; else return erfi;}
}
}
 return 0; //did not converge
}