// test using Multi-dim (2D)  Distribution object interface
// and compare results and CPU performances using TF2::GetRandom
//
// run within ROOT (.x unuranMulti2D.cxx+) or pass any extra parameter in the command line to get
// a graphics output

#include "TStopwatch.h"
#include "TUnuran.h"
#include "TUnuranMultiContDist.h"

#include "TH2.h"
#include "TF2.h"
#include "TCanvas.h"
#include "TMath.h"

#include "TRandom3.h"
#include "TSystem.h"
#include "TApplication.h"
#include "TError.h"

#include "Math/Functor.h"

// #include "Math/ProbFunc.h"
// #include "Math/DistFunc.h"

#define _USE_MATH_DEFINES // for Windows
#include <cmath>
#include <iostream>

//#define DEBUG

using std::cout;
using std::endl;

int n = 1000000;

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

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

   double sigma_x = p[0];
   double sigma_y = p[1];
   double rho = p[2];
   double u = x[0] / sigma_x ;
   double v = x[1] / sigma_y ;
   double c = 1 - rho*rho ;
   double result = (1 / (2 * TMath::Pi() * sigma_x * sigma_y * sqrt(c)))
      * exp (-(u * u - 2 * rho * u * v + v * v ) / (2 * c));
   return result;
}

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

   return std::log( gaus2d(x,p) ); // should re-implement it
}



int testUnuran(TUnuran & unr, const std::string & method, const TUnuranMultiContDist & dist, TH2 * h1, const TH2 * href ) {

#ifdef DEBUG
   n = 1000;
#endif


   // init unuran
   bool ret =   unr.Init(dist,method);
   if (!ret) {
      std::cerr << "Error initializing unuran with method " << unr.MethodName() << std::endl;
      return -1;
   }

   h1->Reset();

   // test first the time
   TStopwatch w;

   w.Start();
   double x[2];
   for (int i = 0; i < n; ++i) {
      unr.SampleMulti(x);
      h1->Fill(x[0],x[1]);
      if (method == "gibbs" && i < 100)
         std::cout << x[0] << " , " << x[1] << std::endl;
   }

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

   double prob = href->Chi2Test(h1,"UU");
   double ksprob = href->KolmogorovTest(h1);
   std::cout << "Time using Unuran  " << unr.MethodName() << "   \t=\t " << time << "\tns/call \t\tChi2 Prob = "
        << prob << "\tKS Prob = " << ksprob << std::endl;
   if (prob < 1E-06) {
      std::cout << "Chi2 Test failed ! " << std::endl;
      href->Chi2Test(h1,"UUP"); // print all chi2 test info
      return 1;
   }
   return 0;
}

int testGetRandom(TF2 * f, TH1 * h1, const TH2 * href = nullptr) {


   // test first the time
   h1->Reset();

   TStopwatch w;
   w.Start();
   double x[2] = {0,0};
   for (int i = 0; i < n; ++i) {
      f->GetRandom2(x[0],x[1]);
      h1->Fill(x[0],x[1]);
   }

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


   if (href != nullptr) {
      double prob = href->Chi2Test(h1,"UU");
      std::cout << "Time using TF1::GetRandom()    \t=\t " << time << "\tns/call \t\tChi2 Prob = "<< prob << std::endl;
      if (prob < 1E-06) {
         std::cout << "Chi2 Test failed ! " << std::endl;
         href->Chi2Test(h1,"UUP"); // print all chi2 test info
         return 2;
      }
   }
   else
      std::cout << "Time using TF1::GetRandom()    \t=\t " << time << "\tns/call\n";
   return 0;
}


int unuranMulti2D() {

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

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


   gSystem->Load("libMathCore");
   gSystem->Load("libUnuran");

   // simple test of unuran



   TH2D * h1 = new TH2D("h1","UNURAN gaussian 2D distribution",100,-10,10,100,-10,10);
   TH2D * h2 = new TH2D("h2","TF1::GetRandom gaussian 2D distribution",100,-10,10,100,-10,10);

   TH2D * h3 = new TH2D("h3","UNURAN truncated gaussian 2D distribution",100,-1,1,100,-1,1);
   TH2D * h4 = new TH2D("h4","TF1::GetRandom truncated gaussian 2D distribution",100,-1,1,100,-1,1);


   TF2 * f = new TF2("g2d",gaus2d,-10,10,-10,10,3);
   double par[3] = {1,1,0.5};
   f->SetParameters(par);

   TF2 * flog = new TF2("logg2d",log_gaus2d,-10,10,-10,10,3);
   flog->SetParameters(par);


   f->SetNpx(100);
   f->SetNpy(100);
   std::cout << " Nimber of function points in TF1, Npx = " << f->GetNpx() << " Npy =  "  << f->GetNpy() << std::endl;


   std::cout << "Test using an undefined domain :\n\n";

   // test first with getrandom
   testGetRandom(f,h2);

   // create multi-dim distribution
   TUnuranMultiContDist dist(f);

   // use directly function interfaces
   //ROOT::Math::Functor f2( *f, 2);
   //TUnuranMultiContDist dist(f2);

   TUnuran unr(gRandom,2);  // 2 is debug level

   int iret = 0;
   TH2 * href = h2;

   //vnrou method (requires only pdf)
   std::string method = "vnrou";
   iret |= testUnuran(unr, method, dist, h1, href);


   //hitro method (requires only pdf)
   method = "hitro";
   iret |= testUnuran(unr, method, dist, h1, href);

   //gibbs requires log of pdf and derivative
//#define USE_GIBBS
#ifdef USE_GIBBS
   method = "gibbs";
   // need to create a new  multi-dim distribution with log of pdf
   TUnuranMultiContDist logdist(flog,0,true);
   iret |= testUnuran(unr, method, logdist, h1, href);
#endif

   // test setting the mode
   std::cout << "\nTest setting the mode in Unuran distribution:\n\n";
   double m[2] = {0,0};
   dist.SetMode(m);

   method = "vnrou";
   iret |= testUnuran(unr, method, dist, h1, href);

   method = "hitro";
   iret |= testUnuran(unr, method, dist, h1, href);

#ifdef USE_GIBBS
   method = "gibbs";
   logdist.SetMode(m);
   iret |= testUnuran(unr, method, logdist, h1, href);
#endif


//    std::cout << " chi2 test of histogram generated with Unuran vs histogram generated with TF1::GetRandom " << std::endl;
//    h1->Chi2Test(h2,"UUP");

//#ifdef LATER

   double xmin[2] = { -1, -1 };
   double xmax[2] = { 1, 1 };

   f->SetRange(xmin[0],xmin[1],xmax[0],xmax[1]);
   // change function domain (not yet implemented in unuran for multidim functions)
   dist.SetDomain(xmin,xmax);

   const double *xlow = dist.GetLowerDomain();
   const double *xup = dist.GetUpperDomain();
   std::cout << "\nTest truncated distribution in domain [ " << xlow[0] << " : " << xup[0]
        << " , " << xlow[1] << " : " << xup[1] << " ] :\n\n";


   testGetRandom(f,h4);

   method = "vnrou";
   iret |= testUnuran(unr, method, dist, h3, h4);
   method = "hitro";
   iret |= testUnuran(unr, method, dist, h3, h4);
#ifdef USE_GIBBS
   logdist.SetDomain(xmin,xmax);
   method = "gibbs";
   iret |= testUnuran(unr, method, logdist, h3, h4);
#endif

//#ifdef LATER
   TCanvas * c1 = new TCanvas("c1_unuran2D","Multidimensional distribution",10,10,900,900);
   c1->Divide(2,2);
#ifdef NO_TRUNC
   TCanvas * c1 = new TCanvas("c1_unuran2D","Multidimensional distribution",10,10,500,500);
   c1->Divide(1,2);
#endif


   c1->cd(1);
   h1->Draw("col");


   c1->cd(2);
   h2->Draw("col");

   c1->cd(3);
   h3->Draw("col");
   c1->cd(4);
   h4->Draw("col");
//#endif

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

   return iret;

}

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

   return iret;
}
#endif