#include "Generator.h"
#include <iostream>
#include "TH1.h"
#include "TH2.h"
#include "TRandom3.h"
#include "TGraph.h"
#include "TMath.h"
#include "TCanvas.h"
#include "Math/DistFunc.h"
#include "TLegend.h"
#include "TFile.h"
#include "TROOT.h"
#include "TStyle.h"
#include "TMatrixD.h"
#include "TMatrixDSym.h"
#include "TVectorD.h"

void Statistics()
{
  // Exercise 2. Check the mean, variance, skweness and kurtosis
  TRandom3* r = new TRandom3();
  Double_t rho = 3./4.;
  TH1F* h = new TH1F("h","",30,0.,1.);
  h->SetTitle("Muon decay");
  h->GetXaxis()->SetTitle("#frac{E}{E_{max}}");
  h->GetYaxis()->SetTitle("Events");

  for (int k = 0; k < 1000000; ++k){
    if (k % 1000 == 0)
      cout << "Producing " << k << "event" << endl;
    h->Fill(Generator(rho,r));
  }
  h->Draw();
  cout << " Mean:     " << h->GetMean()                        << endl;
  cout << " Variance: " << (h->GetStdDev())*(h->GetStdDev())   << endl;
  cout << " Skewness: " << h->GetSkewness()                    << endl;
  cout << " Kurtosis: " << h->GetKurtosis()                    << endl;
  
  
}


Double_t getMean(Double_t rho, TRandom3* r)
{

  TH1F* h     = new TH1F("h","",30,0.,1.);
  
  for (int k = 0; k < 100000; ++k){
    h->Fill(Generator(rho,r));
  }
  
  Double_t mean = h -> GetMean();
  delete h;
  return mean;

}


void Dependency()
// Exercise 3
{
  TRandom3* r = new TRandom3();
  const int nPoints = 20;
  Double_t yval[nPoints];
  Double_t xval[nPoints];

  for (int k = 0; k < nPoints; ++k){
    Double_t rho = double(k)/double(nPoints);
    yval[k] = getMean(rho,r);
    xval[k] = rho;
  }

  TGraph* gr = new TGraph(nPoints,xval,yval);
  gr->SetMarkerStyle(kCircle);
  gr->Draw("AP");
  TF1* straight = new TF1("straight","[0] + x*[1]",0,1);
  gr->Fit(straight);
  gr->GetXaxis()->SetTitle("#rho");
  gr->GetYaxis()->SetTitle("Mean");
  gr->SetTitle("");
  cout << "Independent term: " << straight->GetParameter(0) << endl;
  cout << "Slope:            " << straight->GetParameter(1) << endl; 
  return;

}

void Exercise4()
{
  TRandom3* r = new TRandom3();
  Double_t rho   = 3./4.;
  const int nEvents = 1000;
  const int nExperiments = 5000; // (for the cross-check)

  Double_t xbar = 0.;
  for (int ev = 0; ev < nExperiments; ++ev){
    xbar += Generator(rho,r);
  }
  xbar /= Double_t(nExperiments);
  cout << "Estimator of the mean   " << xbar << endl;
  cout << "Estimator of rho        " << 15.*xbar/2. - 9./2. << endl;
  cout << "Calculation of Var[rho] " << 15.*15.*0.0433007/(2.*2.*float(nEvents)) << endl;

  cout << " /////////////////////////// " << endl;
  cout << " This is the cross check     " << endl;
  cout << " /////////////////////////// " << endl;

  TH1F* hRho  = new TH1F("hRho","",100,0.,1.);
  TH1F* hMean = new TH1F("hMean","",100,0.,1.);
  for (int exp = 0; exp < nExperiments; ++exp){
    xbar = 0;
    for (int ev = 0; ev < nEvents; ++ev){
      xbar += Generator(rho,r);
    }
    xbar /= nEvents; 
    hMean->Fill(xbar);
    hRho-> Fill(15.*xbar/2. - 9./2.);
  }

  TCanvas* c = new TCanvas();
  c -> Divide(1,2);
  c -> cd(1);
  hMean -> Draw();
  c -> cd(2);
  hRho -> Draw();
  cout << "Rho deviation    " << hRho -> GetStdDev() *  hRho -> GetStdDev()  << endl;
  
  return;
}

void Exercise5() 
{
  const int N = 1000; // maximum number of events
  Double_t rho0 = 3./4;
  Double_t rho[N-1];
  Double_t mean[N-1];
  Double_t Number[N-1];
  
  TRandom3* r = new TRandom3();
  for (int k = 1; k < N; k++){
    cout << "Running  "<< k << " experiments" << endl;
    TH1F* h = new TH1F("h","",20,0.,1.);
    for (int exp = 0; exp < k; ++exp){
      h->Fill(Generator(rho0,r));
    }
    mean[k-1] = h->GetMean();
    rho[k-1]  = 15.*h->GetMean()/2. - 9./2.;
    Number[k-1] = k;
    delete h;
  }
  TGraph* gr1 = new TGraph(N-1,Number,rho);
  TGraph* gr2 = new TGraph(N-1,Number,mean);
  TF1*    fun1 = new TF1("fun1","0*x+[0]",0,N);
  TF1*    fun2 = new TF1("fun2","0*x+[0]",0,N);
  fun1->SetParameter(0,rho0);
  cout << 3./5.+2.*rho0/15. << endl;
  fun2->SetParameter(0,3./5.+2.*rho0/15.);
  gr1->GetXaxis()->SetTitle("N events");
  gr1->GetYaxis()->SetTitle("#hat{#rho}");
  gr1->SetTitle("");
  gr2->SetLineColor(kBlue);
  gr2->SetTitle("");
  gr2->GetXaxis()->SetTitle("N events");
  gr2->GetYaxis()->SetTitle("#hat{#mu}");

  TCanvas* c = new TCanvas();
  c -> Divide(1,2);
  c -> cd(1);
  gr1->Draw("AL");
  fun1->Draw("same");
  c -> cd(2);
  gr2->Draw("ALsame");
  fun2->Draw("same");
  return;
}

void Exercise6(int N = 50)
{
  gStyle->SetOptStat(0);
  TRandom3* r = new TRandom3();
  const Double_t rho0 = 3./4.;
  Double_t mu0 = 3./5. + 2*rho0/15.; // true mean value

  TH1F* hStudent = new TH1F("hStudent","",100,-10.,10.);
  const int TRIES = 10000;
  for (int exp = 0; exp < TRIES; ++exp){
    Double_t sum  = 0;
    Double_t sum2 = 0;
    for (int k = 0; k < N; ++k){
      Double_t result = Generator(rho0,r);
      Double_t rhoHat = 15.*result/2. - 9./2.;
      sum  += rhoHat;
      sum2 += rhoHat*rhoHat;
    }
    Double_t s    = sum2/N - (sum/N)*(sum/N);
    s *= double(N)/double(N+1);
    s  = TMath::Sqrt(s);
    Double_t xbar = sum/N;
    Double_t t = (xbar - rho0)/(s/TMath::Sqrt(N));

    
    hStudent->Fill(t,1./double(TRIES)/hStudent->GetBinWidth(1));
  }
  TCanvas* c = new TCanvas();
  TF1 *studentf = new TF1("studentf","TMath::Gamma(([0]+1.0)/2.0)/TMath::Gamma([0]/2.0)/TMath::Sqrt([0]*TMath::Pi())*TMath::Power(1+x*x/[0],-([0]+1.0)/2.)",-10.5,10.5);
  hStudent->Draw();
  hStudent->GetXaxis()->SetTitle("t");
  hStudent->GetYaxis()->SetTitle("Normalized");
  studentf->SetParameter(0,Double_t(N-1));
  studentf->Draw("same");
  c -> SaveAs(Form("Exercise6N_%d.pdf",N));
  return;
}



void Exercise6_run()
// this is not necesarry is just for my own amusement
{
  cout << "[W]: Running Exercise6_run(). This is not part of the problem set. Run Exercise6() instead" << endl;
  Exercise6(10);
  Exercise6(20);
  Exercise6(30);
  Exercise6(50);
  Exercise6(100);

}

void Exercise7()
{
  const double rho = 3./4.;
  TRandom3* r = new TRandom3();

  const int N = 100;
  const int nExp = 10000;
  TH1F* h2 = new TH1F("h2","",50,-0.5,49.5); // number of events in a bin
  
  for (int exp = 0; exp < nExp; ++exp){
    TH1F* h  = new TH1F("h", "",50,0.,1.); // histograms for experiments
     for (int ev =0; ev < N; ++ev){
       h->Fill(Generator(rho,r));
     }
     h2->Fill(h->GetBinContent(40),1./nExp);
     delete h;
  }
  h2->Draw();
  
  Double_t x[50];
  Double_t y[50];
  
  for (int j = 0; j < 50; ++j){
    x[j] = double(j);
    y[j] = ROOT::Math::binomial_pdf(j,3.558/100.,100);
  }
    
  TGraph* gr = new TGraph(50,x,y);
  gr->SetMarkerStyle(kFullCircle);
  gr->Draw("psame");
  gr->SetMarkerColor(kBlack);
  
  h2->SetLineColor(kBlue);

  TLegend* leg = new TLegend(0.661891,0.59448,0.823782,0.781316);
  leg->AddEntry(gr,"Binomial","p");
  leg->AddEntry(h2,"N Entries","l");
  leg->Draw("same");
  return;
}

void Exercise7_pois()
{
  const double rho = 3./4.;
  TRandom3* r = new TRandom3();

  const int N = 300;
  const int nExp = 10000;
  TH1F* h2 = new TH1F("h2","",50,-0.5,49.5); // number of events in a bin
  
  for (int exp = 0; exp < nExp; ++exp){
    if (exp % 1000 == 0) 
      cout << "Exp " << exp << endl;

    TH1F* h  = new TH1F("h", "",50,0.,1.); // histograms for experiments
     for (int ev =0; ev < N; ++ev){
       h->Fill(Generator(rho,r));
     }
     h2->Fill(h->GetBinContent(40),1./nExp);
     delete h;
  }
  h2->Draw();
  
  Double_t x[50];
  Double_t y[50];
  
  for (int j = 0; j < 50; ++j){
    x[j] = double(j);
    y[j] = ROOT::Math::poisson_pdf(j,3.558*N/100.);
  }
    
  TGraph* gr = new TGraph(50,x,y);
  gr->SetMarkerStyle(kFullCircle);
  gr->Draw("psame");
  gr->SetMarkerColor(kBlack);
  
  h2->SetLineColor(kBlue);

  TLegend* leg = new TLegend(0.661891,0.59448,0.823782,0.781316);
  leg->AddEntry(gr,"Poisson","p");
  leg->AddEntry(h2,"N Entries","l");
  leg->Draw("same");
  return;
}

void Exercise7_2()
{
  const int nExp       = 5000;
  const int nEvtPerExp = 500;
  const int nBins      = 20;

  double rho = 3./4.;
  TRandom3* r = new TRandom3();
  TH1F*     hChi = new TH1F("hChi","",50.,0.,3.5*nBins);
  TF1* pdf = new TF1("fun","4*x*x*(3-3*x) + [0]*(8*x*x*(4*x-3)/3)",0,1);
  pdf->SetParameter(0,rho);

  for (int exp = 0; exp < nExp; ++exp){
    if (exp%100 == 0) cout << "Experiment " << exp << endl;
    TH1F* hExp = new TH1F("h","",nBins,0.,1.);

    for (int ev = 0; ev < nEvtPerExp; ++ev){
      hExp -> Fill(Generator(rho,r));
    }

    Double_t sum = 0;
    for (int bin =0; bin < nBins; ++bin){
      Double_t prob    = hExp->GetBinWidth(bin+1)*pdf->Eval(hExp->GetBinCenter(bin+1));  // probability for a event to fall in the bin 
      Double_t central = prob*nEvtPerExp;
      Double_t desv    = nEvtPerExp*prob*(1-prob);
      // sum += ( hExp->GetBinContent(bin+1) - central)*( hExp->GetBinContent(bin+1) - central )/(prob*nEvtPerExp);
      sum += ( hExp->GetBinContent(bin+1) - central)*( hExp->GetBinContent(bin+1) - central )/(prob*nEvtPerExp*(1-prob));
    }

    hChi->Fill(sum,1/double(nExp)/hChi->GetBinWidth(1));
    delete hExp;
  }
  hChi->Draw();
  TF1 *chi2 = new TF1("pdf", "ROOT::Math::chisquared_pdf(x, [0])",0,100);
  chi2 -> SetParameter(0,nBins);
  // chi2 -> Draw("same");
  hChi->Fit(chi2);
  return;
}


void Exercise7_multi()
{
  const int nBins      = 20;
  const int nExp       = 100000;
  const int nEvtPerExp = 1000;
  double rho = 3./4.;
  TRandom3* r = new TRandom3();

  TH1F* hChi      = new TH1F("chi","",50,0.,3*nBins);
  TH2F* hCovMatr  = new TH2F("hCovMatr","",nBins,-0.5,nBins-0.5,nBins,-0.5,nBins-0.5);
  TH2F* hCorrMatr = new TH2F("hCorrMatr","",nBins,-0.5,nBins-0.5,nBins,-0.5,nBins-0.5);
  TF1*  pdf       = new TF1("fun","4*x*x*(3-3*x) + [0]*(8*x*x*(4*x-3)/3)",0,1);
  pdf->SetParameter(0,rho);
  TH1F* hDumm = new TH1F("h","",nBins,0.,1.);

  TMatrixDSym covMatr (nBins);


  for (int binx = 0; binx < hCovMatr->GetXaxis()->GetNbins(); ++binx){
    for (int biny = 0; biny < hCovMatr->GetYaxis()->GetNbins(); ++biny){
      float px = pdf->Eval(hDumm->GetBinCenter(binx+1)) *  hDumm->GetBinWidth(binx + 1);
      float py = pdf->Eval(hDumm->GetBinCenter(biny+1)) *  hDumm->GetBinWidth(biny + 1);
      if (binx == biny){
	hCovMatr->SetBinContent(binx+1,biny+1,nEvtPerExp*px*(1-px));
	hCorrMatr->SetBinContent(binx+1,biny+1,1);
	covMatr(binx,biny) = nEvtPerExp*px*(1-px);
      }
      else{
	hCovMatr->SetBinContent(binx+1,biny+1,-nEvtPerExp*px*py);
	hCorrMatr->SetBinContent(binx+1,biny+1,-TMath::Sqrt(px*py/((1-px)*(1-py))));
	covMatr(binx,biny) = -TMath::Sqrt(px*py/((1-px)*(1-py)));
      }
    }
  }


  TCanvas* c1 = new TCanvas();
  c1 -> Divide(1,2);
  c1 -> cd(1);
  hCovMatr->Draw("colz text");
  c1 -> cd(2);
  hCorrMatr->Draw("colz text");

  covMatr.Invert();



  for (int exp = 0; exp < nExp; ++exp){
    if (exp%100 == 0) cout << "Experiment " << exp << endl;
    TH1F* hExp2 = new TH1F("h2","",nBins,0.,1.);

    for (int ev = 0; ev < nEvtPerExp; ++ev){
      hExp2 -> Fill(Generator(rho,r));
    }

    TVectorD results(nBins);
    for (int bin =0; bin < nBins; ++bin){
      float p = pdf->Eval(hExp2->GetBinCenter(bin+1)) *  hExp2->GetBinWidth(bin + 1);
      results(bin) = hExp2->GetBinContent(bin+1) - p*nEvtPerExp;
    }
    float sum = results*(covMatr*results);
    hChi->Fill(sum,1/double(nExp)/hChi->GetBinWidth(1));
    delete hExp2;
  }

  TCanvas* c2 = new TCanvas();
  hChi->Draw();
  TF1 *chi2 = new TF1("pdf", "ROOT::Math::chisquared_pdf(x, [0])",0,100);
  chi2 -> SetParameter(0,nBins);
  chi2 -> Draw("same");
  c2->SaveAs("result.pdf");
  TFile* _file = TFile::Open("result.root","recreate");
  chi2->Write();
  _file->Close();
}

Double_t Likelihood( Double_t* x, Double_t* par)
{
  // TFile* _file = TFile::Open("tmp.root");
  TH1F* h = (TH1F*) gROOT -> FindObject("h");
  TF1* pdf = new TF1("fun","4*x*x*(3-3*x) + [0]*(8*x*x*(4*x-3)/3)",0,1);
  Double_t prob = 0;
  pdf -> SetParameter(0,x[0]);
  for (int bin = 0; bin < h->GetNbinsX(); ++bin){
    // prob *= 10*ROOT::Math::binomial_pdf(h->GetBinContent(bin+1),pdf->Eval(h->GetBinCenter(bin+1))*h->GetBinWidth(bin+1),par[0]);
    prob += TMath::Log(TMath::Power(pdf->Eval(h->GetBinCenter(bin+1))*h->GetBinWidth(bin+1),h->GetBinContent(bin+1))/TMath::Factorial(h->GetBinContent(bin+1)));
  }
  prob += h->GetEntries()*TMath::Log(h->GetEntries()) - h->GetEntries();
  // cout << prob << endl;

  delete pdf;
  return TMath::Exp(prob);
}

void Exercise8()
{
  TH1F* h          = new TH1F("h","",10,0.,1.);
  TRandom3* r      = new TRandom3();
  const double rho = 3./4.;
  const int nEvtPerExp = 500;

  for (int ev = 0; ev < nEvtPerExp; ++ev){
    h->Fill(Generator(rho,r));
  }
  h->Draw();
  return;


  cout << "Working fine" << endl;
  TF1* likelihood = new TF1("likelihood",Likelihood,0.,1.,1);
  likelihood->SetParameter(0,nEvtPerExp);
  likelihood->Draw();
  cout << "mean                         " << 15.*h->GetMean()/2. - 9./2. << endl;
  cout << " XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX" << endl;
  cout << " XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX" << endl;
  cout << " CUIDADO! EL LIKELIHOOD ESTA MULTIPLICADO POR 10^50 PARA QUE EL FIT CONVERJA" << endl;
  cout << " XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX" << endl;
  cout << " XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX" << endl;
  cout << "Maximum likelihood estimator " << likelihood->GetMaximumX(0.,1.) << endl;
   
}