#include <TROOT.h>
#include <TCanvas.h>
#include <TH1D.h>
#include <TFile.h>
#include <TKey.h>
#include <TMath.h>
#include <TTree.h>
#include <algorithm>
#include <iostream>
#include <fstream>

void NGenAuAuYt(const Int_t nevents, const Double_t yup, const Double_t meff, const Double_t clustermean, const Int_t nclustermax, const Double_t decaymean, const Int_t ndecaymax, Double_t boostmax) {
  
  if (!gROOT->GetClass("TGenPhaseSpace")) gSystem.Load("libPhysics");
  
  // ================================================================
  
  // open output text file
  
  ofstream display;
  TString name = "ngenauauyt.txt";
  display.open(name);
  display << name << endl;
  cout << "// display opened successfully" << endl;
  
  // ================================================================
  
  TStopwatch *tstop = new TStopwatch();
  tstop->Start();
  
  // ================================================================
  
  // write run parameters
  
  display << "    nevents = " << nevents << endl;      // # events in simulation
  display << "        yup = " << yup << endl;          // rapidity histogram upper bound
  display << "       meff = " << meff << endl;         // effective mass of cluster
  display << "clustermean = " << clustermean << endl;  // mean # clusters in event (conditional Poisson dist'n, 0 < n <= nclustermax)
  display << "nclustermax = " << nclustermax << endl;  // maximum # clusters in event 
  display << "  decaymean = " << decaymean << endl;   // mean # of decays in cluster(conditional Poisson dist'n, 1 < n <= ndecaymax))
  display << "  ndecaymax = " << ndecaymax << endl;   // maximum # of decays in cluster
  display << "   boostmax = " << boostmax << endl;    // maximum radial rapidity of cluster
  
  // ================================================================
  
  // create a ROOT file for simulation output
  
  TFile *treeout = new TFile("ngenauauyt.root","RECREATE","ROOT file with NGenAuAuYt tree");
  if ( treeout->IsOpen() ) cout << "// treeout opened successfully" << endl; 

  treeout->cd();
  
  const Int_t _kMaxParticles = 200*nclustermax*ndecaymax;
  const Int_t _kMaxClusters  = 200*nclustermax;

  Double_t yt,phib,pt,phi,yc,px,py,pz,e2,ee,m2,pxsum,pysum,pzsum,eesum,yy,theta,eta,absEta;
  Int_t _eventIndex;
  //Int_t charge, pid;
  //Int_t _nParticlesRead;
  Int_t _nParticipants;
  Int_t _nClusters;
  Int_t _nParticles;
  //Int_t _nBinary;
  //Int_t _nPart1;
  //Int_t _nPart2;
  Int_t _Mult05Eta;
  Int_t _Mult10Eta;
  Int_t _Mult20Eta;
  Int_t _nDecays[_kMaxClusters];
  Double_t _yc[_kMaxClusters];
  Double_t _yt[_kMaxClusters];
  Double_t _phib[_kMaxClusters];
  //Int_t _pid[_kMaxParticles];
  //Int_t _charge[_kMaxParticles];
  Double_t _px[_kMaxParticles];
  Double_t _py[_kMaxParticles];
  Double_t _pz[_kMaxParticles];
  Double_t _ee[_kMaxParticles];
  Double_t _yy[_kMaxParticles];
  Double_t _pt[_kMaxParticles];
  Double_t _phi[_kMaxParticles];
  Double_t _eta[_kMaxParticles];
  //Double_t _m2[_kMaxParticles];
  
  TTree *tree = new TTree("tree","ROOT tree with NGenAuAuYt branches");
  
  tree->Branch("eventIndex",      &_eventIndex,      "_eventIndex/I");
  tree->Branch("nParticipants",   &_nParticipants,   "_nParticipants/I");
  tree->Branch("nClusters",       &_nClusters,       "_nClusters/I");
  tree->Branch("nParticles",      &_nParticles,      "_nParticles/I");
  //tree->Branch("nBinary",         &_nBinary,         "_nBinary/I");
  //tree->Branch("nPart1",          &_nPart1,          "_nPart1/I");
  //tree->Branch("nPart2",          &_nPart2,          "_nPart2/I");
  tree->Branch("Mult05Eta",  &_Mult05Eta, "_Mult05Eta/I");
  tree->Branch("Mult10Eta",  &_Mult10Eta,  "_Mult10Eta/I");
  tree->Branch("Mult20Eta",  &_Mult20Eta,  "_Mult20Eta/I");
  tree->Branch("yc",              _yc,               "_yc[_nClusters]/D");
  tree->Branch("yt",              _yt,               "_yt[_nClusters]/D");
  tree->Branch("phib",            _phib,             "_phib[_nClusters]/D");
  tree->Branch("nDecays",         _nDecays,          "_nDecays[_nClusters]/I");
  //tree->Branch("pid",             _pid,              "_pid[_nParticles]/I");
  //tree->Branch("charge",          _charge,           "_charge[_nParticles]/I");
  //tree->Branch("px",              _px,               "_px[_nParticles]/D");
  //tree->Branch("py",              _py,               "_py[_nParticles]/D");
  tree->Branch("pz",              _pz,               "_pz[_nParticles]/D");
  tree->Branch("ee",              _ee,               "_ee[_nParticles]/D");
  tree->Branch("yy",              _yy,               "_yy[_nParticles]/D");
  tree->Branch("pt",              _pt,               "_pt[_nParticles]/D");
  tree->Branch("phi",             _phi,              "_phi[_nParticles]/D");
  tree->Branch("eta",             _eta,              "_eta[_nParticles]/D");
  //tree->Branch("m2",              _m2,               "_m2[_nParticles]/D");

  tree->SetMaxTreeSize(1000000);
  
  //================================================================
  
  const Double_t Pi = 3.14159265358979;
  
  // histograms to monitor decay dist'ns
  
  TH1D *hParticipants;  
  TH1D *hClusters;
  TH1D *hDecays;
  TH1D *hRapidity;
  TH1D *hBoost;
  TH1D *hPhiB;
  
  Int_t nbinsy = Int_t(yup)*20;
  Int_t nbinsboost = Int_t(boostmax*10.);
  
  hParticipants = new TH1D("hParticipants","dNpart/dx",400,0.,400.);  
  hClusters = new TH1D("hClusters","dNclus/dx (Poisson)",nclustermax + 1,0.,Double_t(nclustermax) + 1.); 
  hDecays = new TH1D("hDecays","dNclus/dNdecay (Poisson)",ndecaymax + 1,0.,Double_t(ndecaymax) + 1.); // limited by TGenPhaseSpace to 18 decays
  hRapidity = new TH1D("hRapidity","dNclus/dy (Uniform)",nbinsy,-yup,yup); 
  hBoost = new TH1D("hBoost","dNclus/dyt (Linear)",nbinsboost,0.,boostmax); 
  hPhiB = new TH1D("hPhiB","dNclus/dphi (Uniform)",36,-Pi,Pi); 
  
  // ================================================================
  
  // random variables for simulation  
  
  TFile *f; 
  TH1D *dNpdx;
  
  f = new TFile("dNpdx.root", "READ");
  
  dNpdx = (TH1D*)f->Get("fdNpdx");
  cout << "// got Npart density" << endl;
  
  TRandom2 *genCluster = new TRandom2();
  TRandom2 *genParticles = new TRandom2();
  TRandom2 *genRapidity = new TRandom2();
  TRandom2 *genBoost = new TRandom2();
  TRandom2 *genPhiB = new TRandom2();
  
  TF1 *f1 = new TF1("f1","x",0.,boostmax);
  
  Int_t npart,ncollisions,nclusters,ntracks,ndecays,ntracks; 
  Double_t genPart; 
  
  // ================================================================
  
  // begin event loop
  
  for (Int_t n = 0; n < nevents; n++) {
    
    if (n < 2) display << "" << endl;    
    if (n < 2) display << "Event " << n+1 << endl;
    
    _eventIndex    = n;
    _nParticipants = 0;
    _nClusters     = 0;
    _nParticles    = 0;
    _Mult05Eta     = 0;
    _Mult10Eta     = 0;
    _Mult20Eta     = 0;
    
    //========================================================
    
    // generate ncollisions for event using Ant's AuAu Run 4 NPart histogram
    
    genPart = fdNpdx->GetRandom();
    npart = Int_t(genPart);
    
    _nParticipants = npart;
    
    hParticipants->Fill(Double_t(npart));
    
    ncollisions = Int_t(genPart/2.);
    
    if (n < 2) display << "Collisions = " << ncollisions << endl;
    
    //========================================================
    
    // begin ncollisions loop
    
    for (Int_t ncol = 0; ncol < ncollisions; ncol++) {
      
      // rv feedback:
      if (n < 2) display << "" << endl;      
      if (n < 2) display << "   Collision " << ncol+1 << endl;
      
      //========================================================
      
      // generate nclusters for collision with Poisson nclustermax, clustermean
      
      nclusters = 0;
      
      while (nclusters < 1 || nclusters > nclustermax) {
	nclusters = genCluster->Poisson(clustermean);
	
	if (nclusters > nclustermax) display << "   nclusters = " << nclusters << endl;
	
      }  
      
      hClusters->Fill(Double_t(nclusters));  
      
      if (n < 2) display << "   Clusters = " << nclusters << endl;
      
      //========================================================
      
      // generate transverse rapidity for all clusters in collision;
      // generate radial coordinates uniformly on a disc
      // boost with transverse rapidity proportional to r
      
      yt = f1->GetRandom();
      
      hBoost->Fill(yt);
      _yt[_nClusters] = yt;
      
      phib = genPhiB->Uniform(-Pi,Pi);
      
      hPhiB->Fill(phib);
      _phib[_nClusters] = phib;
      
      pt = meff*sinh(yt);
      
      px = pt*cos(phib);
      py = pt*sin(phib); 
      
      //========================================================
      
      // begin ncluster loop
      
      for (Int_t nc = 0; nc < nclusters; nc++) {
	
	if (n < 2) display << "" << endl;      
	if (n < 2) display << "      Cluster " << nc+1 << endl;
	
	// generate longitudinal rapidity for cluster
	
	yc = genRapidity->Uniform(-yup,yup);
	
	hRapidity->Fill(yc);
	_yc[_nClusters] = yc;
	
	// generate 4-vector for cluster
	
	pz = meff*sinh(yc);
	
	e2 = meff*meff + px*px + py*py + pz*pz;
	ee = sqrt(e2);
	
	TLorentzVector W(px,py,pz,ee);	
	
	//========================================================
	
	// generate ndecays for the cluster with Poisson decaymean, ndecaymax
	
	ndecays = 0;
	
	while (ndecays < 2 || ndecays > ndecaymax) {
	  ndecays = genParticles->Poisson(decaymean);
	  
	  //if (ndecays > ndecaymax) display << "" << endl;
	  //if (ndecays > ndecaymax) display << "event " << _eventIndex << ", ndecays = " << ndecays << endl;
	  
	}   
	
	if (n < 2) display << "      Decays = " << ndecays << endl;
	
	hDecays->Fill(Double_t(ndecays));  
	
	_nDecays[_nClusters] = ndecays;
	_nClusters++;

	if (n < 2) display << "" << endl;
	if (n < 2) display << "         yt = " << yt << ", phib = " << phib << ", yc = " <<  yc << ", meff = " << meff << endl;
	if (n < 2) display << "         Ptot,M2 = (" << W.Px() << "," << W.Py() << "," << W.Pz() << "," << W.E() << "," << W.M2() << ")" << endl;
	if (n < 2) display << "" << endl;    
	
	//========================================================
	
	// assign masses to cluster decays
	
	Double_t *mass = new Double_t[ndecays];
	
	TLorentzVector *pPi = new TLorentzVector[ndecays];
	
	for (Int_t p = 0; p < ndecays; p++) {
	  mass[p] = 0.140;
	  
	  if (n < 2) display << "         mass[" << p << "] = " << mass[p] << endl;
	  
	}
	
	if (n < 2) display << "" << endl;
	
	// decay cluster into ndecays
	
	TGenPhaseSpace cluster;
	cluster.SetDecay(W, ndecays, mass);
	
	Double_t weight = cluster.Generate();
	
	// accounting for cluster 4-vector Ptot
	
	pxsum = 0.;
	pysum = 0.;
	pzsum = 0.;
	eesum = 0.;
	
	// begin ndecay loop
	
	for (Int_t p = 0; p < ndecays; p++) {
	  
	  pPi[p] = cluster.GetDecay(p);
	  
	  px = pPi[p].Px();
	  py = pPi[p].Py();
	  pz = pPi[p].Pz();
	  ee = pPi[p].E();
	  m2 = pPi[p].M2();
	  
	  if (n < 2) display << "         P" << p+1 << " = (" << pPi[p].Px() << "," << pPi[p].Py() << "," << pPi[p].Pz() << "," << pPi[p].E() << "," << pPi[p].M2() << ")" << endl;
	  
	  // cluster accounting
	  
	  pxsum += px;
	  pysum += py;
	  pzsum += pz;
	  eesum += ee;
	  
	  // variables for tree branches	  
	  
	  px = pPi[p].Px();
	  py = pPi[p].Py();
	  pz = pPi[p].Pz();
	  ee = pPi[p].E();
	  m2 = pPi[p].M2();
	  
	  yy = .5*log( (ee+pz)/(ee-pz) );
	  
	  pt  = sqrt(px*px+py*py);
	  phi = atan2(py,px);
	  //if (pt>0.001) {
	  theta  = atan2(pt, pz);
	  //}else{
	  //  theta  = atan2(0.001, pz);
	  //}
	  eta    = -log(tan(0.5*theta));
	  absEta = fabs(eta);
	  
	  //if (absEta<6.)
	  //  {
	  //_pid[_nParticles]    = pid;
	  //_charge[_nParticles] = charge;
	  //_px[_nParticles]     = px;
	  //_py[_nParticles]     = py;
	  _pz[_nParticles]     = pz;
	  _ee[_nParticles]     = ee;
	  _yy[_nParticles]     = yy;
	  _pt[_nParticles]     = pt;
	  _phi[_nParticles]    = phi;
	  _eta[_nParticles]    = eta;
	  _nParticles++;
	  
	  //if(absEta <= 2.0  && charge!=0)
	  if(absEta <= 2.0)
	    {
	      _Mult20Eta++;  
	      if(absEta <= 1.0)
		{
		  _Mult10Eta++;
		  if(absEta <= 0.5) _Mult05Eta++;
		}
	    }
	  //} //absEta<6.	
	  
	}
	
	if ( fabs(pxsum) < .0000000000001 ) pxsum = 0.; 
	if ( fabs(pysum) < .0000000000001 ) pysum = 0.; 
	if ( fabs(pzsum) < .0000000000001 ) pzsum = 0.; 
	if ( fabs(eesum) < .0000000000001 ) eesum = 0.; 
	
	// cluster accounting s/b meff
	
	if (n < 2) display << "" << endl;
	if (n < 2) display << "         Event " << n+1 <<  "/" << nevents << ", Collision " << ncol+1 << "/" << ncollisions << ", Cluster " << nc+1 << "/" << nclusters << ":" << endl;
	if (n < 2) display << "         Ptot,'M2' = (" << pxsum << "," <<  pysum << "," << pzsum << "," << eesum << "," << eesum*eesum - pzsum*pzsum - pysum*pysum - pxsum*pxsum << ")" << endl;

	delete [] pPi;
	delete [] mass;
	
      }  // end cluster loop
      
    }  // end collision loop
    
    // fill the tree
    
    tree->Fill();   
    
    //=====================================================
    
    // event counter
    
    //if (n % 10000 == 0) display << "events = " << n << endl;
    if (n % 1000 == 0) cout << "events = " << n << endl;
    if (n % 1000 == 0) tstop->Stop();
    if (n % 1000 == 0) tstop->Print();
    if (n % 1000 == 0) cout << "runtime = " << tstop->RealTime() << endl;
    if (n % 1000 == 0) cout << "CPUtime = " << tstop->CpuTime() << endl;
    if (n % 1000 == 0) tstop->Start(kFALSE);
    
    //==========================================================
  
    
  } // end event loop
  
  display << "// event loop complete" << endl;
  cout << "// event loop complete" << endl;

  //=======================================================
  
  // print tree to out file
  
  tree->Print();

  // save all objects in this file

  tree->Write();
  
  tree->Delete();
  
  treeout->Close();
  
  cout << "// treeout closed" << endl;
    
  //========================================================
  
  // cluster, decay and rapidity densities
  
  TH1D *fParticipants;
  TH1D *fClusters;
  TH1D *fDecays;
  TH1D *fRapidity;
  TH1D *fBoost;
  TH1D *fPhiB;
  TString name,title;
  
  Double_t sum,binsize;
  
  sum = hParticipants->Integral();
  
  name = "fParticipants";	    
  title = "dNpart/dx density";
  fParticipants = (TH1D*)hParticipants->Clone();
  fParticipants->SetName(name);
  fParticipants->SetTitle(title);
  fParticipants->Scale( 1./sum );
  
  sum = hClusters->Integral();
  
  name = "fClusters";	    
  title = "dNclus/dx density";
  fClusters = (TH1D*)hClusters->Clone();
  fClusters->SetName(name);
  fClusters->SetTitle(title);
  fClusters->Scale( 1./sum );
  
  sum = hDecays->Integral();
  
  name = "fDecays";	    
  title = "dNclus/dNdecay density";
  fDecays = (TH1D*)hDecays->Clone();
  fDecays->SetName(name);
  fDecays->SetTitle(title);
  fDecays->Scale( 1./sum );
  
  sum = hRapidity->Integral();
  
  binsize = 2.*yup/Double_t(nbinsy);
  
  name = "fRapidity";	    
  title = "dNev/dy density";
  fRapidity = (TH1D*)hRapidity->Clone();
  fRapidity->SetName(name);
  fRapidity->SetTitle(title);
  fRapidity->Scale( 1./(sum*binsize) );
  
  sum = hBoost->Integral();
  
  binsize = boostmax/Double_t(nbinsboost);
  
  name = "fBoost";	    
  title = "dNev/dy density";
  fBoost = (TH1D*)hBoost->Clone();
  fBoost->SetName(name);
  fBoost->SetTitle(title);
  fBoost->Scale( 1./(sum*binsize) ); 
  
  sum = hPhiB->Integral();
  
  binsize = 2.*Pi/36.;
  
  name = "fPhiB";	    
  title = "dNev/dy density";
  fPhiB = (TH1D*)hPhiB->Clone();
  fPhiB->SetName(name);
  fPhiB->SetTitle(title);
  fPhiB->Scale( 1./(sum*binsize) );  
  
  
  // ================================================================
  
  // calculate probabilities for dNclus/dx density
  
  Double_t p0,p1,pk,pksum,kpksum,pktail;
  Double_t *pclus = new Double_t[nclustermax];
  Double_t *pdecay = new Double_t[ndecaymax-1];
  Double_t *apclus = new Double_t[nclustermax+2];
  Double_t *apdecay = new Double_t[ndecaymax+2];
  
  pksum = 0.;
  kpksum = 0.;
  pktail = 0.;
  
  display << "" << endl;
  display << "Calculated Poisson(" << clustermean << ") probabilities:" << endl;
  
  for (Int_t k = 0; k <= nclustermax+10; k++) {
    pk = TMath::PoissonI(Double_t(k),clustermean);  // pk is p(k)
    if (k > nclustermax) pktail += pk;
    display << "p(" << k << ") = " << pk << endl;
  }
  
  display << "" << endl;
  display << "Conditional dNclus/dx probabilities:" << endl;
  
  p0 = TMath::PoissonI(0.,clustermean);
  
  display << "   p0 = " << p0 << endl;
  display << "pktail = " << pktail << endl;
  display << "measure = " << 1. - p0 - pktail << endl;
  
  for (Int_t k = 1; k <= nclustermax; k++) {
    pk = TMath::PoissonI(Double_t(k),clustermean); 
    pclus[k-1] = pk/(1.-p0-pktail);                // pclus(k-1) is conditional p(k)
    pksum += pclus[k-1];
    kpksum += Double_t(k)*pclus[k-1]; 
    display << "conditional p(" << k << ") = " << pclus[k-1] << endl;
  }
  
  display << "   total probability = " << pksum << endl;
  display << "mean of distribution = " << kpksum << endl; 
  
  // ================================================================
  
  // actual dNclus/dx distributions
  
  display << "" << endl;
  display << "Actual dNclus/dx probabilities:" << endl; 
  
  for (Int_t k = 1; k <= nclustermax+2; k++) {
    apclus[k-1] = fClusters->GetBinContent(k);
    display << "conditional p(" << k-1 << ") = " << apclus[k-1] << endl;
  }
  
  display << "   total probability = " << fClusters->Integral() << endl;
  display << "mean of distribution = " << fClusters->GetMean() << endl;  
  
  // ================================================================
  
  // calculate probabilities for dNclus/dNdecay density
  
  display << "" << endl;
  display << "Calculated Poisson(" << decaymean << ") probabilities:" << endl;
  
  pksum = 0.;
  kpksum = 0.;
  pktail = 0.;
  
  for (Int_t k = 0; k <= ndecaymax+10; k++) {
    pk = TMath::PoissonI(Double_t(k),decaymean); 
    if (k > ndecaymax) pktail += pk;
    display << "p(" << k << ") = " << pk << endl;
  }
  
  display << "" << endl;
  display << "Conditional dNclus/dNdecay probabilities:" << endl;
  
  p0 = TMath::PoissonI(0.,decaymean);
  p1 = TMath::PoissonI(1.,decaymean);
  
  display << "   p0 = " << p0 << endl;
  display << "   p1 = " << p1 << endl;
  display << "pktail = " << pktail << endl;
  display << "measure = " << 1. - p0 - p1 - pktail << endl;
  
  for (Int_t k = 2; k <= ndecaymax; k++) {
    pk = TMath::PoissonI(Double_t(k),decaymean); 
    pdecay[k-2] = pk/(1.-p0-p1-pktail);
    pksum += pdecay[k-2];
    kpksum += Double_t(k)*pdecay[k-2]; 
    display << "conditional p(" << k << ") = " << pdecay[k-2] << endl;
  }
  
  display << "   total probability = " << pksum << endl;
  display << "mean of distribution = " << kpksum << endl; 
  
  // ================================================================
  
  // actual dNclus/dNdecay distributions
  
  display << "" << endl;
  display << "Actual dNclus/dNdecay probabilities:" << endl; 
  
  for (Int_t k = 1; k <= ndecaymax+2; k++) {
    apdecay[k-1] = fDecays->GetBinContent(k);
    display << "conditional p(" << k-1 << ") = " << apdecay[k-1] << endl;
  }
  
  display << "   total probability = " << fDecays->Integral() << endl;
  display << "mean of distribution = " << fDecays->GetMean() << endl; 
  
  //==========================================================
  
  display.close(); 
  
  cout << "// display closed" << endl;
  
  TFile *histout = new TFile("ngenauauyth.root","RECREATE","ROOT file with NGenAuAuYt histograms");
  if ( histout->IsOpen() ) cout << "// histout opened successfully" << endl; 

  histout->cd();
  
  hParticipants->Write();
  hClusters->Write();
  hDecays->Write();
  hRapidity->Write();
  hBoost->Write();
  hPhiB->Write();
  
  fParticipants->Write();
  fClusters->Write();
  fDecays->Write();
  fRapidity->Write(); 
  fBoost->Write();
  fPhiB->Write();
  
  //=======================================================
  
  histout->Close();
  
  cout << "// histout closed" << endl;
  
  //=======================================================
  
  cout << "events = " << n << endl;
  tstop->Stop();
  tstop->Print();
  cout << "runtime = " << tstop->RealTime() << endl;
  cout << "CPUtime = " << tstop->CpuTime() << endl;
  
  //==========================================================
  
}