// Assume we have a bias t in the probability of the outcome of
// the die throw: 
// p(1) = 1/6-t/2;
// p(2)... p(5) = 1/6-t/8;
// p(6) = 1/6+t;
// Let us fit for t with a explicit calculation of the max likelihood

#include <iostream>
#include <sstream>
#include <fstream>

#include "TProfile.h"
#include "TH1.h"
#include "TF1.h"
#include "TH1D.h"
#include "TH2.h"
#include "TStyle.h"
#include "TCanvas.h" 
#include "TMath.h"
#include "TROOT.h"
#include "Riostream.h"
#include "TMath.h"
#include "TRandom.h"
#include "TRandom3.h"
using namespace std;

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


void Die5(int N=100, int Nrep=100, double alpha=0.05) {
  
  // Alpha is the type-1 error rate we admit. We accept the alternative hypothesis if the
  // probability that we observe data at least as discrepant with the null hypothesis (no bias in the die)
  // is smaller than alpha.
  double nsigma = sqrt(2)*ErfInverse(1.-alpha); // sqrt(2)*TMath::ErfInverse(1.-alpha);

  TH1D * Powernum = new TH1D ("Powernum","",100,-0.1666667,0.3333333);
  TH1D * Powerden = new TH1D ("Powerden","",100,-0.1666667,0.3333333);
  TH1D * Coverage = new TH1D ("Coverage","",100,-0.1666667,0.3333333);
  TH2D * Values = new TH2D ("Values","",100,-0.16666667,0.3333333,100,-0.17,0.34);

  for (int it=0; it<100; it++) {
    double ttrue = -1./6. + (double)it/200. +0.0025;
    double ok=0;
    int Nrepok=0;

    for (int rep=0; rep<Nrep; rep++) {
      int n1=0;
      int n2=0;
      int n3=0;
      for (int i=0; i<N; i++) {
	double r = gRandom->Uniform(0.,1.);
	if (r<1./6.-ttrue/2.) {
	  n1++;
	} else if (r>1.-1./6.-ttrue) {
	  n3++;
	} else {
	  n2++;
	}
      }
      /*   cout << "Die throwing" << endl; */
      /*   cout << "Results:" << endl; */
      /*   cout << "n1 = " << n1 << endl; */
      /*   cout << "n2 = " << n2 << endl; */
      /*   cout << "n3 = " << n3 << endl; */
      /*   cout << endl; */
      
      // quadratic equation for max of L: coefficients
      double a = 18*(n1+n2+n3);
      double b = -3*(7*n1+n2+10*n3);
      double c = -(4*n1+n2-8*n3);
      
      double t1,t2;
      double discr=b*b-4*a*c;
      double tmin=-1./6.;
      double tmax=1./3.;
      double tstar=0;
      double rms;
      bool onesol=false;
      if (discr<0) {
	cout << "No solution for max likelihood" << endl;
	// This never happens in the considered case
      } else {
	t1 = (-b-sqrt(discr))/(2*a);
	t2 = (-b+sqrt(discr))/(2*a);
	if (t1>=tmin && t1<=tmax) {
	  if (t2>=tmin && t2<=tmax) {
	    // NNBB when n1=0 or n3=0 there is always one solution 
	    // at the boundary
	    if (t1-tmin>tmax-t2) {
	      //cout << "Bias is estimated to be t = " << t1;
	      tstar=t1;
	    } else {
	      //cout << "Bias is estimated to be t = " << t2;
	      tstar=t2;
	    }
	  } else {
	    //cout << "Bias is estimated to be t = " << t1;
	    tstar=t1;
	    onesol=true;
	  }
	} else if (t2>=tmin && t2<=tmax) {
	  //cout << "Bias is estimated to be t = " << t2;
	  tstar=t2;
	  onesol=true;
	}
	// determine error from inverse of second derivative of logL
	double d2logl;
	if (tstar>-1/6. && tstar<1/3.) {
	  d2logl=9*n1/pow(1-3*tstar,2)+
	    9*n2/pow(4-3*tstar,2)+
	    36*n3/pow(1+6*tstar,2);
	} else {
	  d2logl=100000.;
	}
	rms = sqrt(1/d2logl);
	//cout << " +- " << rms << endl;
	Nrepok++;
	if (fabs(tstar-ttrue)<rms) ok++;
	Values->Fill(ttrue,tstar);

	if (fabs(tstar/rms)>nsigma) Powernum->Fill(ttrue);
	Powerden->Fill(ttrue);

	// turn into p-values
	/*   cout << endl; */
	/*   cout << "This means the following probabilities:" << endl; */
	/*   cout << "p(1) = " << 1./6.-tstar/2. << " +- " << rms/2. << endl; */
	/*   cout << "p(x) = " << 1./6.+tstar/8. << " +- " << rms/8. << endl; */
	/*   cout << "p(6) = " << 1./6.+tstar << " +- " << rms << endl; */
      } // end if discr>0     
    } // end Nrep loop
    double ratio=ok/(double)Nrepok;
    Coverage->Fill(ttrue,ratio);
    Coverage->SetBinError(it+1,sqrt(ratio*(1-ratio)/(double)Nrep));
    //cout << "In total the error bar on t covers " << ok/(double)Nrep*100 << "% of the times" << endl;
  }
  Powernum->Divide(Powerden);

  TCanvas *C = new TCanvas ("C","",600,600);
  C->Divide(1,3);
  C->cd(1);
  Coverage->SetLineColor(kRed);
  Coverage->SetMinimum(0);
  Coverage->SetMaximum(1.1);
  Coverage->Fit("pol0","","",-0.15,0.3);
  C->cd(2);
  Values->Draw();
  C->cd(3);
  Powernum->SetMinimum(0);
  Powernum->SetMaximum(1.1);
  Powernum->Draw();
}


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
}