#include <string.h>
#include <iostream>
#include <vector>
#include "TFile.h"
#include "TLorentzVector.h"
#include "TH1.h"
#include "TH2.h"
#include "TF1.h"
#include "TRandom3.h"
#include "TVector3.h"
#include <fstream>

using namespace std;
 
struct myGen_t {
  int runNumber;
  int eventNumber;
  double weight;
  vector<int> pidVecGeant;
  vector<int> pidVecPDG;
  vector<int> chargeVec;
  vector<TLorentzVector> mo4Vec;
  vector<TVector3> vertex3Vec;
};


//FUNCTION PROTOTYPES
void printUsage(); // print usage
void writeEventLund(myGen_t myGen,std::ofstream *lundFile);
void writeEventGenr8(myGen_t myGen,std::ofstream *genr8File);


//MAIN
int main(int argc, char **argv) {
  
  char  inFileName[120];
  char  outFileName[120];
  char  outFileNameLund[120];
  char  outFileNameGenr8[120];
  char *argptr;
  int nToRun = 100000;
  int rSeed = 0;
  int outType  = 1;
  int runNumber = 99999;
  //Set the default output file
  sprintf(outFileName,"./myOutFile.root");
  sprintf(outFileNameLund,"./myOutFile.lund");
  sprintf(outFileNameGenr8,"./myOutFile.genr8");
  //if (argc == 1) printUsage();
  for (int i=1; i<argc; i++) {
    argptr = argv[i];
    if (*argptr == '-') {
      argptr++;
      switch (*argptr) {
      case 'h':
        printUsage();
        break;
      case 'n':
        nToRun = atoi(++argptr);
        break;
      case 't':
        outType = atoi(++argptr);
        break;
      case 'r':
        runNumber = atoi(++argptr);
        break;
      case 's':
        rSeed = atoi(++argptr);
        break;
      case 'o':
        strcpy(outFileName,++argptr);
        strcpy(outFileNameLund,++argptr);
        strcpy(outFileNameGenr8,++argptr);
        break;
      default:
        printUsage();
        break;
      }
    } else {
      strcpy(inFileName,argptr);
    }
  }

  cout<<"outType = "<<outType<<endl;

  //Create histograms here
  TH1D *hDummy1d     = new TH1D("hDummy1d","",100,0.0,1.0);
  TH2D *hDummy2d     = new TH2D("hDummy2d","",100,0.0,1.0,100,0.0,1.0);
  
  TH1D * hEGamma = new TH1D("hEGamma","",200,7.5,9.5);
  TH1D * hTVal = new TH1D("hTVal","",200,0.0,20);
  TH2D * hMoVsThetaPro = new TH2D("hMoVsThetaPro","",90,0.0,90.0,100,0.0,10.0);
  TH2D * hMoVsThetaRho = new TH2D("hMoVsThetaRho","",90,0.0,90.0,100,0.0,10.0);
  TH1D * hMassRho = new TH1D("hMassRho","",200,0.0,2.0);
  TH1D * hMassPipPim = new TH1D("hMassPipPim","",200,0.0,2.0);
  TH2D * hMoVsThetaPim = new TH2D("hMoVsThetaPim","",90,0.0,90.0,100,0.0,10.0);
  TH2D * hMoVsThetaPip = new TH2D("hMoVsThetaPip","",90,0.0,90.0,100,0.0,10.0);

  //Setup the output files
  std::ofstream lundFile;
  std::ofstream genr8File;

  if (outType == 2) lundFile.open(outFileNameLund);
  if (outType == 3) genr8File.open(outFileNameGenr8);
  

  gRandom =  new TRandom3();
  gRandom->SetSeed(rSeed);


  double pi = 3.14159;
  
  double mPro       = 0.938;
  double mPip       = 0.13957;
  double mRhoCenter = 0.7755;
  double rhoWidth   = 0.1464;
  double mRho       = mRhoCenter;
  
  TLorentzVector target(0,0,0,mPro);
    
  double eGammaLow  = 8.0;
  double eGammaHigh = 9.0;  
  
  for(Int_t i=0;i<nToRun;i++){
    if (i%10000 == 0 && i > 0) cout<<"Events processed = "<<i<<endl;
    bool keep = false;
    double eGamma = 0;
    double testVal = 0;
    while (keep == false){
      eGamma = gRandom->Uniform(eGammaLow,eGammaHigh);
      testVal = gRandom->Uniform(0.0,1.0/eGammaLow);
      if (testVal < 1.0/eGamma) keep = true;
    }
    hEGamma->Fill(eGamma);
    
    TLorentzVector beam(0,0,eGamma,eGamma);
    TLorentzVector pTot4Vec = beam + target;
    double sVal = pTot4Vec.Mag2();
    double wVal = pTot4Vec.Mag();  //Invariant mass of initial state
    //
    //beam and target to cm frame
    //
    TLorentzVector beamCM = beam;
    beamCM.Boost(-pTot4Vec.BoostVector());
    TLorentzVector targetCM = target;
    targetCM.Boost(-pTot4Vec.BoostVector());
    
    keep = false;
    while (keep == false){
      mRho = gRandom->BreitWigner(mRhoCenter,rhoWidth);
      if (mRho >= 2*mPip && wVal >= (mRho+mPro)) keep = true;
    }
    hMassRho->Fill(mRho);
    
    TLorentzVector pRhoCM4Vec;
    TLorentzVector pProCM4Vec;
    
    //Since wVal is invariant, it is the same in every frame
    //The center of mass frame is especially nice. Let us "decay"
    //the initial state into the final state = proton and rho meson
    //
    double eRhoCM = (sVal + pow(mRho,2) - pow(mPro,2))/(2*wVal);
    double pRhoCM = sqrt(pow(eRhoCM,2) - pow(mRho,2));
    double eProCM = wVal - eRhoCM;
    double pProCM = pRhoCM;
    keep = false;
    while (keep == false){
      //
      //Now phase-space generate the angles
      double phiRho      = gRandom->Uniform(-pi,pi);
      double cosThetaRho = gRandom->Uniform(-1.0,1.0);
      //Obtain the components of the rho 4-vector
      double pRhoZCM          = pRhoCM*cosThetaRho;
      double pRhoTransverseCM = sqrt(pow(pRhoCM,2) - pow(pRhoZCM,2));
      double pRhoXCM          = pRhoTransverseCM*cos(phiRho);
      double pRhoYCM          = pRhoTransverseCM*sin(phiRho);
      //Obtain the components of the proton 4-vector
      double pProXCM          = -pRhoXCM;
      double pProYCM          = -pRhoYCM;
      double pProZCM          = -pRhoZCM;
      //Create the rho CM 4-vector and the proton CM 4-vector
      pRhoCM4Vec = TLorentzVector(pRhoXCM,pRhoYCM,pRhoZCM,eRhoCM);
      pProCM4Vec = TLorentzVector(pProXCM,pProYCM,pProZCM,eProCM);
      //
      //Find extreme values of t (theta = 0 and pi)
      TLorentzVector pProCM4VecExtA = TLorentzVector(0,0,pProCM,eProCM);
      TLorentzVector pProCM4VecExtB = TLorentzVector(0,0,-pProCM,eProCM);
      TLorentzVector p1MinusP3ExtA = targetCM - pProCM4VecExtA;
      TLorentzVector p1MinusP3ExtB = targetCM - pProCM4VecExtB;
      double tValAbsExtB = abs(p1MinusP3ExtB.Mag2());
      TLorentzVector p1MinusP3 = targetCM - pProCM4Vec;
      double tValAbs = abs(p1MinusP3.Mag2());
      double bVal = -0.5;
      double expTest = exp(bVal*tValAbs);
      double expTestMax = exp(bVal*tValAbsExtB);
      double randTest = gRandom->Uniform(0.0,expTestMax);
      if (randTest < expTest) keep = true;
      
    }
    //
    //Boost back to lab
    //
    TLorentzVector pPro4Vec = pProCM4Vec;
    pPro4Vec.Boost(pTot4Vec.BoostVector());
    
    TLorentzVector pRho4Vec = pRhoCM4Vec;
    pRho4Vec.Boost(pTot4Vec.BoostVector());
    
    TLorentzVector p1MinusP3 = target - pPro4Vec;
    double tValAbs = abs(p1MinusP3.Mag2());
    hTVal->Fill(tValAbs);
    hMoVsThetaPro->Fill(pPro4Vec.Theta()*180.0/3.14159,pPro4Vec.P());
    hMoVsThetaRho->Fill(pRho4Vec.Theta()*180.0/3.14159,pRho4Vec.P());
    
    //////Rho decay////////////
    //Let RFR = Rest frame of Rho
    double ePipRFR = mRho/2.0;
    double pPipRFR = sqrt(pow(ePipRFR,2) - pow(mPip,2));
    double phiPipRFR      = gRandom->Uniform(-pi,pi);
    double cosThetaPipRFR = gRandom->Uniform(-1.0,1.0);
    //Obtain the components of the pip 4-vector
    double pPipZRFR          = pPipRFR*cosThetaPipRFR;
    double pPipTransverseRFR = sqrt(pow(pPipRFR,2) - pow(pPipZRFR,2));
    double pPipXRFR          = pPipTransverseRFR*cos(phiPipRFR);
    double pPipYRFR          = pPipTransverseRFR*sin(phiPipRFR);
    //Load the RFR 4-vectors for pip and pim
    TLorentzVector pPipRFR4Vec = TLorentzVector(pPipXRFR,pPipYRFR,pPipZRFR,ePipRFR);
    TLorentzVector pPimRFR4Vec = TLorentzVector(-pPipXRFR,-pPipYRFR,-pPipZRFR,ePipRFR);
    ///////////////////////////
    //Boost the pions to the lab frame from the Rest-Frame of the Rho (RFR)
    //
    TLorentzVector pPip4Vec = pPipRFR4Vec;
    pPip4Vec.Boost(pRho4Vec.BoostVector());
    
    TLorentzVector pPim4Vec = pPimRFR4Vec;
    pPim4Vec.Boost(pRho4Vec.BoostVector());
    
    TLorentzVector pPipPim4Vec = pPip4Vec + pPim4Vec;
    hMassPipPim->Fill(pPipPim4Vec.Mag());
    
    hMoVsThetaPip->Fill(pPip4Vec.Theta()*180.0/3.14159,pPip4Vec.P());
    hMoVsThetaPim->Fill(pPim4Vec.Theta()*180.0/3.14159,pPim4Vec.P());

    
    //Make a vertex vector for the event
    TVector3 vertVec(0,0,0);
  
    //Load the myGen structure with final-state particles
    myGen_t myGen;
    myGen.weight      = 1.0;
    myGen.runNumber   = runNumber;
    myGen.eventNumber = i+1;
    myGen.vertex3Vec.push_back(vertVec);
    //Proton
    myGen.mo4Vec.push_back(pPro4Vec);
    myGen.pidVecPDG.push_back(2212);
    myGen.pidVecGeant.push_back(14);
    myGen.chargeVec.push_back(1);
    //Pi+
    myGen.mo4Vec.push_back(pPip4Vec);
    myGen.pidVecPDG.push_back(211);
    myGen.pidVecGeant.push_back(8);
    myGen.chargeVec.push_back(1);
    //Pi-
    myGen.mo4Vec.push_back(pPim4Vec);
    myGen.pidVecPDG.push_back(-211);
    myGen.pidVecGeant.push_back(9);
    myGen.chargeVec.push_back(-1);

    if (outType == 2) writeEventLund(myGen, &lundFile);
    if (outType == 3) writeEventGenr8(myGen, &genr8File);
    
    
  }
  if (outType == 1){
    TFile *outFile = new TFile(outFileName,"RECREATE");
    outFile->cd();
    //Write histograms to file
    hDummy1d->Write();
    hDummy2d->Write();

    hEGamma->Write();
    hTVal->Write();
    hMoVsThetaPro->Write();
    hMoVsThetaRho->Write();
    hMassRho->Write();
    hMassPipPim->Write();
    hMoVsThetaPim->Write();
    hMoVsThetaPip->Write();

    outFile->Write();
  }

  if (outType == 2) lundFile.close();
  if (outType == 3) genr8File.close(); 

  return 0;
}


void printUsage(){
  fprintf(stderr,"\nUsage:");
  fprintf(stderr,"\nmyEventGen [-switches]\n");
  fprintf(stderr,"\nSWITCHES:\n");
  fprintf(stderr,"-h\tPrint this message\n");
  fprintf(stderr,"-s<arg>\tRandom seed. Default is 0\n");
  fprintf(stderr,"-o<arg>\tOutFileName. Default is myOutFile.txt\n\n");
  fprintf(stderr,"-r<arg>\tRun number. Default is 99999\n\n");
  fprintf(stderr,"-t<arg>\tOutFile type: 1->Root, 2->Lund, 3->Genr8\n\n");

  cout<<"The current default operation is equivalent to the command:"<<endl;
  cout<<"myEventGen -s0 -t1 -r99999 -n100000 -omyOutFile.root\n"<<endl;

  exit(0);
}

void writeEventLund(myGen_t myGen, std::ofstream *lundFile){
  //See https://gemc.jlab.org/gemc/html/documentation/generator/lund.html
  //LUND HEADER
  *lundFile<<myGen.pidVecPDG.size();//  1) Number of particles
  *lundFile<<"\t0.938";             //  2) Mass of target:          NOT REQUIRED
  *lundFile<<"\t1";                 //  3) Atomic number of target: NOT REQUIRED
  *lundFile<<"\t0";                 //  4) Target polarization:     NOT REQUIRED
  *lundFile<<"\t0";                 //  5) Beam polarization
  *lundFile<<"\t11";                //  6) Beam type:               NOT REQUIRED
  *lundFile<<"\t10";                //  7) Beam energy:             NOT REQUIRED
  *lundFile<<"\t2212";              //  8) Interacted nucleon ID:   NOT REQUIRED
  *lundFile<<"\t0";                 //  9) Process ID:              NOT REQUIRED
  *lundFile<<"\t"<<myGen.weight;    // 10) Event weight:            NOT REQUIRED
  endl(*lundFile); //Write the header

  //LUND particle
  for (int i = 0; i<(int)myGen.pidVecPDG.size(); i++){
    *lundFile<<i+1;                          //  1) Particle index
    *lundFile<<"\t0";                        //  2) Lifetime:            NOT REQUIRED
    *lundFile<<"\t1";                        //  3) Type (1 is active)
    *lundFile<<"\t"<<myGen.pidVecPDG[i];     //  4) Particle ID
    *lundFile<<"\t0";                        //  5) Index of parent:     NOT REQUIRED
    *lundFile<<"\t-1";                       //  6) Index of daughter:   NOT REQUIRED
    *lundFile<<"\t"<<myGen.mo4Vec[i].Px();   //  7) Momentum in x (GeV)
    *lundFile<<"\t"<<myGen.mo4Vec[i].Py();   //  8) Momentum in y (GeV)
    *lundFile<<"\t"<<myGen.mo4Vec[i].Pz();   //  9) Momentum in z (GeV)
    *lundFile<<"\t"<<myGen.mo4Vec[i].E();    // 10) Energy (GeV):        NOT REQUIRED
    *lundFile<<"\t"<<myGen.mo4Vec[i].Mag();  // 11) Mass (GeV):          NOT REQUIRED
    *lundFile<<"\t"<<myGen.vertex3Vec[i].X();// 12) Vertex x (cm)
    *lundFile<<"\t"<<myGen.vertex3Vec[i].Y();// 12) Vertex x (cm)
    *lundFile<<"\t"<<myGen.vertex3Vec[i].Z();// 12) Vertex x (cm)
    endl(*lundFile); //Write the particle
  }
}

void writeEventGenr8(myGen_t myGen, std::ofstream *genr8File){
  //See page 5 of https://halldweb.jlab.org/DocDB/0000/000011/001/genr8_v_1_0.pdf
  //Genr8 HEADER
  *genr8File<<myGen.runNumber;                // 1) Run number
  *genr8File<<"\t"<<myGen.eventNumber;        // 2) Event number
  *genr8File<<"\t"<<myGen.pidVecGeant.size(); // 3) Number of particles
  endl(*genr8File); //Write the header

  //Genr8 particle
  for (int i = 0; i<(int)myGen.pidVecGeant.size(); i++){
    *genr8File<<i+1;                          // 1) Particle index
    *genr8File<<"\t"<<myGen.pidVecGeant[i];   // 2) Particle ID
    *genr8File<<"\t"<<myGen.mo4Vec[i].Mag();  // 3) Mass
    endl(*genr8File);                         //Write line
    *genr8File<<"\t"<<myGen.chargeVec[i];     // 1) Charge 
    *genr8File<<"\t"<<myGen.mo4Vec[i].Px();   // 2) Momentum in x
    *genr8File<<"\t"<<myGen.mo4Vec[i].Py();   // 3) Momentum in y
    *genr8File<<"\t"<<myGen.mo4Vec[i].Pz();   // 4) Momentum in z
    *genr8File<<"\t"<<myGen.mo4Vec[i].E();    // 5) Energy
    endl(*genr8File); //Write the particle
  }
}