// test using discrete distribution. 
// Generate numbers from a given probability vector or from a discrete distribution like 
// the Poisson distribution. 
// Compare also the Unuran method for generating Poisson numbers with TRandom::Poisson
//
// run within ROOT (.x unuranDiscrete.cxx+) or pass any extra parameter in the command line to get  
// a graphics output (./unuranDiscrete 1 )


#include "TUnuran.h"
#include "TUnuranDiscrDist.h"

#include "TH1.h"
#include "TMath.h"
#include "TF1.h"
#include "TRandom3.h"
#include "TApplication.h"
#include "TCanvas.h"
#include "TStopwatch.h"
#include "TError.h"


#ifdef WRITE_OUTPUT
#include "TFile.h"
#include "TGraph.h"
#endif

#include "Math/Functor.h"
#include "Math/DistFunc.h"
#include "Math/Util.h"

#include <iostream>
#include <iomanip>

//#define DEBUG
#ifdef DEBUG
int n = 100;
#else
int n = 1000000;
#endif
TCanvas * c1; 
int icanv = 1;

bool useRandomSeed = false;   // to use a random seed different every time

double poisson_pmf(double * x, double * p) { 

   double y = ROOT::Math::poisson_pdf(int(x[0]),p[0]);
//   std::cout << x[0] << " f(x) = " << y << std::endl;
   return y; 
}
double binomial_pmf(double * x, double * p) { 

   double y = ROOT::Math::binomial_pdf(static_cast<unsigned int>(x[0]),p[1], static_cast<unsigned int>(p[0]));
//   std::cout << x[0] << " f(x) = " << y << std::endl;
   return y; 
}

int testUnuran(TUnuran & unr, double & time, TH1 * h1, const TH1 * href,  bool weightHist = false ) { 


   // test first the time
   TStopwatch w; 

   w.Start(); 
   for (int i = 0; i < n; ++i) 
      unr.SampleDiscr(); 

   w.Stop(); 
   time = w.CpuTime()*1.E9/n; 


   // test quality (use cdf to avoid zero bins)
   h1->Reset();
   int n2 = n/10;
   int x;
   for (int i = 0; i<n2; ++i) { 
      x =  unr.SampleDiscr(); 
      h1->Fill( double( x) ); 
   } 
   double prob; 
   if (weightHist) 
      prob = h1->Chi2Test(href,"UW");
   else
      prob = h1->Chi2Test(href,"UU");
   std::string s = "Time using Unuran  " +  unr.MethodName();
   std::cout << std::left << std::setw(40) << s << "\t=\t " << time << "\tns/call \t\tChi2 Prob = "<< prob << std::endl;
   if (prob < 1E-06) { 
      std::cout << "Chi2 Test failed for method " << unr.MethodName() << std::endl;
      if (weightHist) 
         h1->Chi2Test(href,"UWP"); // print all chi2 test info
      else 
         h1->Chi2Test(href,"UUP"); // print all chi2 test info

      return 1;
   }
   return 0; 
}  

int testRootPoisson(double mu, double &time, TH1 * h2) { 
   TStopwatch w; 

   w.Start(); 
   for (int i = 0; i < n; ++i) 
      gRandom->Poisson(mu);

   w.Stop(); 
   time = w.CpuTime()*1.E9/n; 

   // make ref histo
   int n2 = n/10;
   for (int i = 0; i< n2; ++i) { 
      h2->Fill ( gRandom->Poisson(mu) );
   }

   std::string s = "Time using TRandom::Poisson";
   std::cout << std::left << std::setw(40) << s << "\t=\t " << time << std::endl;
   return 0;
}

int testRootBinomial(int m, double p, double &time, TH1 * h2) { 

   TStopwatch w; 

   w.Start(); 
   for (int i = 0; i < n; ++i) 
      gRandom->Binomial(m,p);

   w.Stop(); 
   time = w.CpuTime()*1.E9/n; 

   // make ref histo
   int n2 = n/10;
   for (int i = 0; i< n2; ++i) { 
      h2->Fill ( gRandom->Binomial(m,p) );
   }

   std::string s = "Time using TRandom::Binomial";
   std::cout << std::left << std::setw(40) << s << "\t=\t " << time << std::endl;
   return 0;
}

int testProbVector() { 

   int iret = 0;

   TH1D * h0 = new TH1D("h0","reference prob",10,-0.5,9.5);
   TH1D * h1 = new TH1D("h1","UNURAN gen prob",10,-0.5,9.5);


   double p[10] = {1.,2.,3.,5.,3.,2.,1.,0.5,0.3,0.5 };
   for (int i = 0; i< 10; ++i) { 
      h0->SetBinContent(i+1,p[i]);
      h0->SetBinError(i+1,0.);
   }
   double sum = h0->GetSumOfWeights();
   std::cout << " prob sum = " << sum << std::endl; 
      
   TUnuran unr(gRandom,2); 

   TUnuranDiscrDist dist(p,p+10);

//    TUnuranDiscrDist fDDist = dist; 
//    std::cout << fDDist.ProbVec().size() << std::endl;

   std::cout << "Test  generation with a PV :\n\n";
   
   double time;

   bool ret = unr.Init(dist,"method=dau");
   if (!ret) iret = -1;
   iret |= testUnuran(unr,time,h1,h0,true);

   ret = unr.Init(dist,"method=dgt");
   if (!ret) iret = -1;
   iret |= testUnuran(unr,time,h1,h0,true);

   // dss require the sum
   dist.SetProbSum(sum);
   ret = unr.Init(dist,"method=dss");
   if (!ret) iret = -1;
   iret |= testUnuran(unr,time,h1,h0,true);

   



   c1->cd(icanv++);
   h1->Sumw2();
   h1->Scale( h0->GetSumOfWeights()/(h1->GetSumOfWeights() ) ); 
   h0->Draw();
   h1->SetLineColor(kBlue);
   h1->Draw("Esame");
   c1->Update();

   return iret; 
}

int testPoisson() { 
   // test with a function 
   // Poisson distribution

   int iret = 0;
   std::cout << "\nTest  generation with a Probability function (Poisson) :\n\n";

   TF1 * f = new TF1("f",poisson_pmf,1,0,1);

   // loop on mu values for Nmu times 

   const int Nmu = 5;
   double muVal[Nmu] = {5,10,20,50,100};
   double tR[Nmu]; 
   double tU[Nmu]; 
   double tUdari[Nmu];
   double tUdsrou[Nmu];

  TUnuran unr(gRandom,2); 

  for (int imu = 0; imu < Nmu; ++imu) {  

     const double mu = muVal[imu]; 

     int nmax = static_cast<int>(mu*3);
     TH1D * h2 = new TH1D("h2","reference Poisson prob",nmax,0,nmax);
     TH1D * h3 = new TH1D("h3","UNURAN gen prob",nmax,0,nmax);
     TH1D * h4 = new TH1D("h4","UNURAN gen prob 2",nmax,0,nmax);


     std::cout << "\nPoisson mu = " << mu << std::endl << std::endl;

     testRootPoisson(mu,tR[imu],h2);

     // test UNURAN with standard method
     bool ret = unr.InitPoisson(mu);
     if (!ret) iret = -1;
     testUnuran(unr,tU[imu],h3,h2);

      
     // test changing all the time the mu
     // use re-init for a fast re-initialization 
     TStopwatch w;
     unr.InitPoisson(mu,"dstd");
     double p[1]; p[0] = mu; 
     w.Start(); 
     for (int i = 0; i < n; ++i) {
        unr.ReInitDiscrDist(1,p);  
        int k = unr.SampleDiscr(); 
        if (n % 10 == 0) h4->Fill(k);
     }
     w.Stop(); 
     double time = w.CpuTime()*1.E9/n; 
     double prob = h2->Chi2Test(h4,"UU");
     std::string s = "Time using Unuran w/ re-init  method=" + unr.MethodName(); 
     std::cout << std::left << std::setw(40) << s << "\t=\t " << time 
               << "\tns/call \t\tChi2 Prob = "<< prob << std::endl;
      
     if (prob < 1E-06) { 
        std::cout << "Chi2 Test failed for re-init " << std::endl;
        iret = -2; 
     }
 
     f->SetParameter(0,mu);

#ifdef USE_FUNCTOR
     ROOT::Math::Functor1D f2(f);
     TUnuranDiscrDist dist2 = TUnuranDiscrDist(f2);
#else 
     TUnuranDiscrDist dist2 = TUnuranDiscrDist(f);
#endif

     // dari method (needs mode and pdf sum)
     dist2.SetMode(int(mu) );
     dist2.SetProbSum(1. );
     ret = unr.Init(dist2,"method=dari");
     if (!ret) iret = -1;

     iret |= testUnuran(unr,tUdari[imu],h3,h2);

     // dsrou method (needs mode and sum)

     dist2.SetProbSum(1);
     ret = unr.Init(dist2,"method=dsrou");
     if (!ret) iret = -1;

     iret |= testUnuran(unr,tUdsrou[imu],h3,h2);


     c1->cd(icanv++);
     h2->DrawCopy();
     h3->SetLineColor(kBlue);
     h3->DrawCopy("Esame");
     c1->Update();

     delete h2; 
     delete h3;
     delete h4;
  }

#ifdef WRITE_OUTPUT

   TFile * file = new TFile("unuranPoisson.root","RECREATE");
   // create graphs with results
   TGraph * gR = new TGraph(Nmu,muVal,tR);   
   TGraph * gU = new TGraph(Nmu,muVal,tU);
   TGraph * gU2 = new TGraph(Nmu,muVal,tUdari);
   TGraph * gU3 = new TGraph(Nmu,muVal,tUdsrou);
   gR->Write();
   gU->Write();
   gU2->Write();
   gU3->Write();
   file->Close();

#endif

   return iret; 
}

int testBinomial() { 
   // test using binomial distribution
   int iret = 0; 

   std::cout << "\nTest  generation with a Probability function (Binomimal) :\n\n";

   TF1 * f = new TF1("f",binomial_pmf,1,0,2);

   // loop on binomual values

   const int NBin = 3;
   double pVal[NBin] = {0.5,0.1,0.01};
   double NVal[NBin] = {20,100,1000}; 
   double tR[NBin]; 
   double tU[NBin]; 
   double tUdari[NBin];
   double tUdsrou[NBin];


  TUnuran unr(gRandom,2); 


  for (int ib = 0; ib < NBin; ++ib) {  

    
     double par[2]; 
     par[0] = NVal[ib];
     par[1] = pVal[ib];

     int nmax = static_cast<int>(par[0]*par[1]*3);
     TH1D * h2 = new TH1D("h2","reference Binomial prob",nmax,0,nmax);
     TH1D * h3 = new TH1D("h3","UNURAN gen prob",nmax,0,nmax);
     TH1D * h4 = new TH1D("h4","UNURAN gen prob 2",nmax,0,nmax);


     std::cout << "\nBinomial n = " << par[0] << "  " << par[1] << std::endl;

     testRootBinomial(static_cast<int>(par[0]),par[1],tR[ib],h2);


     // test UNURAN with standard method
     bool ret = unr.InitBinomial(static_cast<int>(par[0]),par[1]);
     if (!ret) iret = -1;
     testUnuran(unr,tU[ib],h3,h2);

      
     // test changing all the time the mu
     // use re-init for a fast re-initialization 

     TStopwatch w;
     unr.InitBinomial(static_cast<int>(par[0]), par[1],"dstd");
     w.Start(); 
     for (int i = 0; i < n; ++i) {
        unr.ReInitDiscrDist(2,par);  
        int k = unr.SampleDiscr(); 
        if (n % 10 == 0) h4->Fill(k);
     }
     w.Stop(); 
     double time = w.CpuTime()*1.E9/n; 
     double prob = h2->Chi2Test(h4,"UU");
     std::string s = "Time using Unuran w/ re-init  method=" + unr.MethodName(); 
     std::cout << std::left << std::setw(40) << s << "\t=\t " << time 
               << "\tns/call \t\tChi2 Prob = "<< prob << std::endl;
      
     if (prob < 1E-06) { 
        std::cout << "Chi2 Test failed for re-init " << std::endl;
        iret = -2; 
     }

     // test the universal methods 

     f->SetParameters(par);
     TUnuranDiscrDist dist = TUnuranDiscrDist(f);

     // dari method (needs mode and pdf sum)
     dist.SetMode(int(par[0]*par[1]) );
     dist.SetProbSum(1. );
     ret = unr.Init(dist,"method=dari");
     if (!ret) iret = -1;

     iret |= testUnuran(unr,tUdari[ib],h3,h2);

     // dsrou method (needs mode and sum)

     ret = unr.Init(dist,"method=dsrou");
     if (!ret) iret = -1;

     iret |= testUnuran(unr,tUdsrou[ib],h3,h2);


     c1->cd(icanv++);
     h2->DrawCopy();
     h3->SetLineColor(kBlue);
     h3->DrawCopy("Esame");
     c1->Update();

     delete h2; 
     delete h3;
     delete h4;
  }


   return iret; 
}

int unuranDiscrete() {

   int iret = 0; 

   c1 = new TCanvas("c1_unuranDiscr_PV","Discrete distribution from PV",10,10,800,800); 
   c1->Divide(3,3);

   // switch off printing of  info messages from chi2 test
   gErrorIgnoreLevel = 1001; 

   // check if using a random seed
   if (useRandomSeed) gRandom->SetSeed(0);

   iret |= testProbVector(); 
   iret |= testPoisson(); 
   iret |= testBinomial(); 

   if (iret != 0) 
      std::cerr <<"\n\nUnuRan Discrete Distribution Test:\t  Failed !!!!!!!" << std::endl;
   else 
      std::cerr << "\n\nUnuRan Discrete Distribution Test:\t OK" << std::endl;
   return iret; 

   return iret;
}


#ifndef __CINT__
int main(int argc, char **argv)
{
   int iret = 0; 
   if (argc > 1) { 
      TApplication theApp("App",&argc,argv);
      iret =  unuranDiscrete();
      theApp.Run();
   } 
   else 
      iret =  unuranDiscrete();
   
}
#endif