#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"


using namespace std;
 
//FUNCTION PROTOTYPES
void printUsage(); // print usage

//MAIN
int main(int argc, char **argv) {
  
  char  inFileName[120];
  char  outFileName[120];
  char *argptr;
  int nToRun = 100000;
  int rSeed = 0;

  //Set the default output file
  sprintf(outFileName,"./myOutFile.root");
  //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 's':
        rSeed = atoi(++argptr);
        break;
      case 'o':
        strcpy(outFileName,++argptr);
        break;
      default:
        printUsage();
        break;
      }
    } else {
      strcpy(inFileName,argptr);
    }
  }

  //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 file                                                       
  TFile *outFile = new TFile(outFileName,"RECREATE");

  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);
      mRho = gRandom->Uniform(1.4,2.5);
      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());
    
  }
  

  outFile->cd();
  //Write histograms to file here
  hDummy1d->Write();
  hDummy2d->Write();

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

  outFile->Write();


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

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

  exit(0);
}