/* bsub -q phenix_cas -R linux -o outlogs/err_kT1_0.log 'runpythia outputs/pp_kT1_0.root 1 M > outputs/pp_kT1_0.log' */
#include <iostream.h>
#include <stdlib.h>
#include <time.h>

#include <TROOT.h>
#include <TH1.h>
#include <TFile.h>
#include <TObject.h>
#include <TNtuple.h>
#include <TPythia6.h>
#include <TMCParticle6.h>
#include <TRandom.h>
#include <TLorentzVector.h>
#include <TTree.h>
#include <TStopwatch.h>

#include "header.h"
#include "pytpart.h"

Double_t pi = TMath::Pi();

int main(int argc, char **argv){

  Int_t PRINTOUT = 0;
  
  TROOT root("pythia","run pythia");

  Int_t   numberEvents;   
  if ( argc<=1 ) {
    cout << "usage: runpythia <outfile> <numberOfEvents> [k/m]" << endl;
    exit(1);
  } else if(argc==3){
    numberEvents = atoi(argv[2]);
  } else {
    numberEvents = atoi(argv[2]);
    if ( argv[3][0]=='k' || argv[3][0]=='K' ) numberEvents *= 1000;
    if ( argv[3][0]=='m' || argv[3][0]=='M' ) numberEvents *= 1000000;
  }

  // ====================
  // CREATE PYTHIA OBJECT
  // ====================
  TPythia6 pythia;

  // ============================================================
  // INITIALIZATION SECTION: SET ANY PYTHIA INPUT PARAMETERS HERE 
  // ============================================================
  // set heavy quark masses:
  pythia.SetPMAS(pythia.Pycomp(4), 1, 1.4); // charm  (KF=4) mass is 1.4GeV
  pythia.SetPMAS(pythia.Pycomp(5), 1, 4.2); // bottom (KF=5) mass is 4.2GeV

  // set primordial kT distribution:
  pythia.SetMSTP(91,1);         // 0 -> no kT;  1 -> gauss; 2 -> exp
  pythia.SetPARP(91,3.);        // set width of kT gaussian: sqrt(<kt^2>)=3GeV/c
  pythia.SetPARP(93,9.);        // upper cut-off for primordial kT distribution (Def=2.0GeV/c)

  // set flavor-dependent fragmentation function (p.365) - Peterson/SLAC for charm&bottom
  pythia.SetMSTJ(11,3);
  // set parton distribution function to CTEQ-5L (p.194)
  pythia.SetMSTJ(51,7);

  // set elementary processes:
  //--
  //  pythia.SetMSEL(0); // take full user control over processes
  // then switch your processes on, one by one, eg.:
  // pythia.SetMSUB(82,1); // g+g->Q+Qbar(heavy quark) - apparently g+g->b+bbar only!
  //  pythia.SetMSUB(53,1);    // g+g->q+qbar, light q+qbar final states eliminated later
  //--
   pythia.SetMSEL(4); //1,2=QCD, c-cbar =4; b-bar=5; Z=11(ISUB=1),W=12(ISUB=2)
   //  pythia->SetMSUB(2,1); // for single W,Z production





  // set the random seed
  Int_t seed = (Int_t)(time(NULL)/3);
  if( (seed>=0) && (seed<=900000000) ) {
    pythia.SetMRPY(1, seed);			// set seed
    pythia.SetMRPY(2, 0);			// use new seed
    cout<<"Random Seed : "<<seed<<endl; 
  } else {cout << "error: time " << seed << " is not valid" << endl; exit(2);}

  // pp collision at 200 GeV:
  Char_t  targ[10] = "p";
  Char_t  proj[10] = "p";
  pythia.Initialize("CMS", targ, proj, 200);



  // ============================================
  // CREATE ROOT OUTPUT FILE AND DEFINE ROOT TREE
  // ============================================
  TClonesArray *particleArray;
  TMCParticle  *particle = new TMCParticle();
  TMCParticle  *parent   = new TMCParticle();
  TMCParticle  *parton1  = new TMCParticle();
  TMCParticle  *parton2  = new TMCParticle();

  header  *hdr = new header;
  pytpart *mypart = new pytpart;
  TClonesArray *part = new TClonesArray("pytpart",1500);
  TFile *outFile = new TFile(argv[1], "RECREATE", "pythia TTree output");
  TTree *pyttree = new TTree("pyttree", "pythia ttree output");
  pyttree->Branch("event", "header", &hdr , 32000, 1);    
  pyttree->Branch("part" , &part, 32000, 1);

  pythia.Pystat(4);  // prints a table of kinem cuts CKIN(I) p.135 and p.251
  TStopwatch timer;
  timer.Start();
  system("date");

  Int_t iPar=0, iFC=0, flavor=-1, status=-1, charge=999, qkf1=0, qkf2=0;
  Double_t Pt=0., Theta=0., Phi=0., E=0., qpt1=0., qpt2=0., qphi1=0., qphi2=0., qthe1=0., qthe2=0.;
  // ==========================
  // START THE MAIN EVENT LOOP
  // ========================== 
  for(Int_t event=1; event<=numberEvents; event++){

    if(event%10000==0) cout<<"Event #"<<event<<"->"<<Double_t(event)/numberEvents*100<<"%"<<endl;
      
    pythia.GenerateEvent();
    
    if (event<=1 && PRINTOUT) {
      // print out particle listing information if desired (this shows decay modes and code numbers)
      pythia.Pylist(1);   // list full pythia generated event
      pythia.Pylist(2);   // list of particle created during the history of the first event
      pythia.Pylist(12);  // list of all decay modes for all defined particles
    }

    // Get list of particle and global event information
    particleArray = (TClonesArray *)pythia.ImportParticles();
    Int_t numParticles = particleArray->GetEntries();        // get number of generated particles

    Int_t nstore = 0;
    // Save final state particles only  ====================================================
    // pythia.Pyedit(11);
    part->Clear();
    particleArray->Clear();
    particleArray = (TClonesArray *)pythia.ImportParticles();
    numParticles  = particleArray->GetEntries();        // get number of generated particles

    int OutPart1 = pythia.GetMSTI(21); // KF flavour codes for outcoming partons to the hard interaction 
    int OutPart2 = pythia.GetMSTI(22);

    // Select only g+g->c+cbar and g+g->b+bar processes:
    //    if(abs(OutPart1)!=4 && abs(OutPart1)!=5) continue;
    //    if(abs(OutPart2)!=4 && abs(OutPart2)!=5) continue;

    int LinePart2 = pythia.GetMSTI(7); // line numbers for documentation of outgoing partons/particles
    int LinePart1 = pythia.GetMSTI(8); // from hard scattering for 2->2 processes (else = 0)
    if(LinePart1>0) {
      parton1 = (TMCParticle *)particleArray->At(LinePart1);
      if(!parton1){cout<<"first parton, line #"<<LinePart1<<", does not exist!"<<endl; return -2;}
      qkf1 = parton1->GetKF();
      if(abs(qkf1)==abs(OutPart1)){
	qpt1  = sqrt(pow(parton1->GetPx(),2)+pow(parton1->GetPy(),2));
	qphi1 = atan2(parton1->GetPy(),parton1->GetPx());
	qthe1 = atan2(Pt,parton1->GetPz());
      } else {
	qpt1  = -9999.;
	qphi1 = -9999.;
	qthe1 = -9999.;
      }
    }
    if(LinePart2>0) {
      parton2 = (TMCParticle *)particleArray->At(LinePart2);
      if(!parton2){cout<<"second parton, line #"<<LinePart2<<", does not exist!"<<endl; return -2;}
      qkf2 = parton2->GetKF();
      if(abs(qkf2)==abs(OutPart2)){
	qpt2  = sqrt(pow(parton2->GetPx(),2)+pow(parton2->GetPy(),2));
	qphi2 = atan2(parton2->GetPy(),parton2->GetPx());
	qthe2 = atan2(Pt,parton2->GetPz());
      } else {
	qpt2  = -9999.;
	qphi2 = -9999.;
	qthe2 = -9999.;
      }
    }

    // now, do the particle loop and store
    for (Int_t i=0; i<numParticles; i++) {
      particle = (TMCParticle *)particleArray->At(i);     // get the particle information

      Pt = sqrt(pow(particle->GetPx(),2)+pow(particle->GetPy(),2));
      Phi = atan2(particle->GetPy(),particle->GetPx());
      Theta = atan2(Pt,particle->GetPz());
      E = particle->GetEnergy();
      flavor = particle->GetKF();
      status = particle->GetKS();
      charge = Int_t(pythia.Pychge(flavor)/3.0);
      
      mypart->SetE(E);
      mypart->SetPt(Pt);
      mypart->SetPhi(Phi);
      mypart->SetTheta(Theta);
      mypart->SetM(particle->GetMass());
      mypart->SetStatus(status);       //Status:"1"=finally or directly produced particles
      mypart->SetFlavor(flavor);       //       "11"=particles which have already decayed           
      mypart->SetCharge(charge);                 

      // PARENT PARTICLE INFORMATION:
      iPar = particle->GetParent();
      if(iPar>0 && iPar<numParticles){
	parent = (TMCParticle *)particleArray->At(iPar);
	if(!parent){ cout<<"ERROR: PARENT particle does not exist!"<<endl; return -1;}
	mypart->SetParent(parent->GetKF());
	mypart->SetEP(parent->GetEnergy());
	mypart->SetPtP(sqrt(pow(parent->GetPx(),2)+pow(parent->GetPy(),2)));
	mypart->SetPhiP(atan2(parent->GetPy(),parent->GetPx()));
	mypart->SetThetaP(atan2(sqrt(pow(parent->GetPx(),2)+pow(parent->GetPy(),2)),parent->GetPz()));
	mypart->SetMP(parent->GetMass());
      } else {
	mypart->SetParent(-9999);
	mypart->SetEP(0.);
	mypart->SetPtP(0.);
	mypart->SetPhiP(0.);
	mypart->SetThetaP(0.);
	mypart->SetMP(0.);
      }
      
      // KEEP TWO CLASSES OF PARTICLES:
      // Status==1 --> final state (childless) particles
      // Status==21 --> initial state partons
      iFC = particle->GetFirstChild();
      if(iFC==0) new ((*part)[nstore++]) pytpart(*mypart);
    } // particle loop
    
    hdr->SetNPART(nstore);                 // Number of particles stored
    hdr->SetX1(pythia.GetPARI(33));        // Bjorken x1,x2 p.158
    hdr->SetX2(pythia.GetPARI(34)); 
    hdr->SetS(pythia.GetPARI(14));      // partonic s,t,u
    hdr->SetT(pythia.GetPARI(15));  
    hdr->SetU(pythia.GetPARI(16)); 
    hdr->SetQ2(pythia.GetPARI(22));      // Q squared
    hdr->SetKF1(qkf1);     
    hdr->SetKF2(qkf2);     
    hdr->SetPT1(qpt1);     
    hdr->SetPT2(qpt2);     
    hdr->SetPHI1(qphi1);     
    hdr->SetPHI2(qphi2);     
    hdr->SetTHE1(qthe1);     
    hdr->SetTHE2(qthe2);     
    
    pyttree->Fill();
    part->Clear();
    particleArray->Clear();
  }

  pythia.Pystat(1);  // prints a table of cross sections etc. p.135 and p.251

  pyttree->Write();
  outFile->Write();
  outFile->Close();
  timer.Stop(); timer.Print();
  return 0;
}