#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);
void twoBodyDecay(TRandom3 *gRandom,double m1,double m2,TLorentzVector wVec,TLorentzVector *p1Vec,TLorentzVector *p2Vec);
void twoBodyDecayTPro(TRandom3 *gRandom,double slopeVal,double mPro,double mR,TLorentzVector wVec,
		      TLorentzVector *pProVec,TLorentzVector *pRVec);


//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;
  int charge = 1;
  //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 'c':
        charge = 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 * hMassR   = new TH1D("hMassR","",200,0.0,2.0);
  TH1D * hMassK892 = new TH1D("hMassK892","",200,0.0,2.0);
  TH1D * hMassKpStar = new TH1D("hMassKpStar","",200,0.0,2.0);
  TH1D * hMassKmStar = new TH1D("hMassKmStar","",200,0.0,2.0);

  TH2D * hKpPi0VsKmPi0 = new TH2D("hKpPi0VsKmPi0","",200,0.0,2.0,200,0.0,2.0);

  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 * 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);
  

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

  
  double mPro       = 0.938272088;
  //double mPip       = 0.13957039;
  double mK         = 0.4937;
  double mKp        = 0.4937;
  double mKm        = 0.4937;
  double mPi0       = 0.1349768;
  
  TLorentzVector target(0,0,0,mPro);
    
  double eGammaLow  = 8.2;
  double eGammaHigh = 8.8;  


  for(Int_t i=0;i<nToRun;i++){
    if (i%10000 == 0 && i > 0) cout<<"Events processed = "<<i<<endl;
    double eGamma = gRandom->Uniform(eGammaLow,eGammaHigh);  
    hEGamma->Fill(eGamma);

    
    TLorentzVector beam(0,0,eGamma,eGamma);
    TLorentzVector pTot4Vec = beam + target;

    double mRLow = 1.36;
    double mRHigh = 1.6;
    double mR = gRandom->Uniform(mRLow,mRHigh);
    bool keep892 = false;
    double mK892 = -1;
    while (keep892 == false){
      mK892 = gRandom->BreitWigner(0.892,0.051);
      //mK892 = 0.892;
      double mRLowTest = mK892 + mK;
      if (mR >= mRLowTest) keep892 = true; 
      if (mK892 <= (mK + mPi0)){
	keep892 = false; 
      
	//cout<<mR<<" : "<<"mK892:mK + mPi0 = "<<mK892<<" : "<<mK+mPi0<<endl;
      }
    }
    //cout<<keep892<<" : "<<"mK892:mK + mPi0 = "<<mK892<<" : "<<mK+mPi0<<endl;
    
    hMassR->Fill(mR);
    hMassK892->Fill(mK892);


    TLorentzVector pProVec,pRVec,pKpVec,pKmVec,pPi0Vec,pK892Vec,pK1Vec,pK2Vec;
    // gamma p -> p R
    //twoBodyDecay(gRandom,mPro,mR,pTot4Vec,&pProVec,&pRVec);
    double tSlope = -1.15;
    twoBodyDecayTPro(gRandom,tSlope,mPro,mR,pTot4Vec,&pProVec,&pRVec);

    //cout<<"pProVec.Mag()  = "<<pProVec.Mag()<<endl;
    //cout<<"pProVec  = "<<pProVec.Px()<<" , "<<pProVec.Py()<<" , "<<pProVec.Pz()<<" , "<<pProVec.E()<<endl;

    //cout<<"pRVec.Mag()  = "<<pRVec.Mag()<<endl;
    //cout<<"pRVec  = "<<pRVec.Px()<<" , "<<pRVec.Py()<<" , "<<pRVec.Pz()<<" , "<<pRVec.E()<<endl;

    // R -> K892 K
    //if (pRVec.Mag() < mK892 + mK) cout<<"pRVec.Mag() < mK892 + mK"<<endl;
    twoBodyDecay(gRandom,mK892,mK,pRVec,&pK892Vec,&pK1Vec);
    // K892 -> K pi0
    twoBodyDecay(gRandom,mK,mPi0,pK892Vec,&pK2Vec,&pPi0Vec);

    //Flip between KpStar and KmStar
    if (charge == -1){
      pKpVec = pK1Vec;
      pKmVec = pK2Vec;
      hMassKmStar->Fill(pK892Vec.Mag());
    } else {
      hMassKpStar->Fill(pK892Vec.Mag());
      pKpVec = pK2Vec;
      pKmVec = pK1Vec;
    } 

    TLorentzVector pProDiffVec =  pProVec - target;
    double tVal = pProDiffVec.Mag2();
    //cout<<"tVal = "<<tVal<<endl;
    hTVal->Fill(fabs(tVal));
 
    TLorentzVector pKpPi0Vec = pKpVec + pPi0Vec;
    TLorentzVector pKmPi0Vec = pKmVec + pPi0Vec;
    hKpPi0VsKmPi0->Fill(pKmPi0Vec.Mag(),pKpPi0Vec.Mag());

    //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(pProVec);
    myGen.pidVecPDG.push_back(2212);
    myGen.pidVecGeant.push_back(14);
    myGen.chargeVec.push_back(1);
    //K+
    myGen.mo4Vec.push_back(pKpVec);
    myGen.pidVecPDG.push_back(321);
    myGen.pidVecGeant.push_back(11);
    myGen.chargeVec.push_back(1);
    //K-
    myGen.mo4Vec.push_back(pKmVec);
    myGen.pidVecPDG.push_back(-321);
    myGen.pidVecGeant.push_back(12);
    myGen.chargeVec.push_back(-1);
    //Pi0
    myGen.mo4Vec.push_back(pPi0Vec);
    myGen.pidVecPDG.push_back(111);
    myGen.pidVecGeant.push_back(7);
    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();
    hMassR->Write();
    hMassK892->Write();
    hKpPi0VsKmPi0->Write();
    //hMassPipPim->Write();
    //hMoVsThetaPim->Write();
    //hMoVsThetaPip->Write();

    hMassKpStar->Write();
    hMassKmStar->Write();

    outFile->Write();
  }

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

  return 0;
}

void twoBodyDecayTPro(TRandom3 *gRandom,double slopeVal,double mPro,double mR,TLorentzVector wVec,
		      TLorentzVector *pPro,TLorentzVector *pR){
  double pi = 3.14159;
  TLorentzVector pRCM4Vec;
  TLorentzVector pProCM4Vec;
  double wVal = wVec.Mag();
  double sVal = wVec.Mag2();
  TLorentzVector pTot4Vec = wVec;
  TLorentzVector targetCM(0,0,0,mPro);
  targetCM.Boost(-pTot4Vec.BoostVector());

  //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 eRCM = (sVal + pow(mR,2) - pow(mPro,2))/(2*wVal);
  double pRCM = sqrt(pow(eRCM,2) - pow(mR,2));
  double eProCM = wVal - eRCM;
  double pProCM = pRCM;
  bool keep = false;
  while (keep == false){
    //                                                                        
    //Now phase-space generate the angles                                     
    double phiR      = gRandom->Uniform(-pi,pi);
    double cosThetaR = gRandom->Uniform(-1.0,1.0);
    //Obtain the components of the R 4-vector                               
    double pRZCM          = pRCM*cosThetaR;
    double pRTransverseCM = sqrt(pow(pRCM,2) - pow(pRZCM,2));
    double pRXCM          = pRTransverseCM*cos(phiR);
    double pRYCM          = pRTransverseCM*sin(phiR);
    //Obtain the components of the proton 4-vector                            
    double pProXCM          = -pRXCM;
    double pProYCM          = -pRYCM;
    double pProZCM          = -pRZCM;
    //Create the R CM 4-vector and the proton CM 4-vector                   
    pRCM4Vec = TLorentzVector(pRXCM,pRYCM,pRZCM,eRCM);
    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 = slopeVal;
    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 pR4Vec = pRCM4Vec;
  pR4Vec.Boost(pTot4Vec.BoostVector());

  *pPro = pPro4Vec;
  *pR = pR4Vec;

}

void twoBodyDecay(TRandom3 *gRandom, double m1,double m2, TLorentzVector wVec, TLorentzVector *p1Vec,TLorentzVector *p2Vec){
  double pi = 3.14159;

  //Let RF = Rest frame
  double sVal = wVec.Mag2();
  double wVal = wVec.Mag();
  double e1RF = (sVal + pow(m1,2) - pow(m2,2))/(2*wVal);
  double e2RF = wVal - e1RF;
  double p1RF = sqrt(pow(e1RF,2) - pow(m1,2));
  double phi1RF      = gRandom->Uniform(-pi,pi);
  double cosTheta1RF = gRandom->Uniform(-1.0,1.0);
  //Obtain the components of the m1 4-vector
  double p1ZRF          = p1RF*cosTheta1RF;
  double p1TransverseRF = sqrt(pow(p1RF,2) - pow(p1ZRF,2));
  double p1XRF          = p1TransverseRF*cos(phi1RF);
  double p1YRF          = p1TransverseRF*sin(phi1RF);
  //Load the RFR 4-vectors for m1 and m2
  TLorentzVector p1RF4Vec = TLorentzVector(p1XRF,p1YRF,p1ZRF,e1RF);
  TLorentzVector p2RF4Vec = TLorentzVector(-p1XRF,-p1YRF,-p1ZRF,e2RF);
  ///////////////////////////
  //Boost back 
  //
  *p1Vec = p1RF4Vec;
  p1Vec->Boost(wVec.BoostVector());
  
  *p2Vec = p2RF4Vec;
  p2Vec->Boost(wVec.BoostVector());
  
}

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
  }
}