#include <iostream>
#include <vector>
using namespace std;
//ROOT libraries
  #include "TTree.h"
  #include "TFile.h"
  #include "TLorentzVector.h"
  #include "TVector3.h"
  #include "TH1.h"
  #include "TH2.h"
  #include "TF1.h"

//My include file          
  #include "myTree.h"



//Structure "Bucket Of Information"
struct myBOI_t {
  TLorentzVector mo;
  int pid;
  int sciDet;
  TVector3 vert;
  int calLayer;
  double calLu;
  double calLv;
  double calLw;
  int nphe; 
  double mass;
  double beta;
  double p;
};
//function prototype
void printUsage();



int main(int argc, char **argv) {
  cout<<"Hello World"<<endl;
  char inFileName[120];
  char outFileName[120];
//Define integer nToRun;
  int nToRun = -1;
  char *argptr;
  int mcData = false;
  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 'o':
        strcpy(outFileName,++argptr);
        break;
      case 'm':
	mcData = true;
        break;

      default:
        printUsage();
        break;
      }
    } else {
      strcpy(inFileName,argptr);
    }
  }
//Define the histograms                                               
  TH2D *hBetaVsP     = new TH2D("hBetaVsP","myTitle",60,0.0,6.0,110,0.0,1.1);
  TH2D *hBetaVsPPro  = new TH2D("hBetaVsPPro","myTitle",60,0.0,6.0,110,0.0,1.1);

  TH1 *hmasspropim =  new TH1D("hmasspropim","Lambda Mass",200,1,1.2);
  TH1 *hmasspropim2 =  new TH1D("hmasspropim2","Lambda Mass",200,1,1.2);


//Setup the output file                                                   
  TFile *outFile = new TFile(outFileName,"RECREATE");


  //Setup the input file and tree                                          
  TFile* inFile = new TFile(inFileName);
  cout<<"inFileName = "<<inFileName<<endl;
  //Get tree out of the file
  TTree *inTree = (TTree*)inFile->Get("clas12");
  
  //Define the tree structure
  myTree_t myTree;
  
  //Get the tree branches into myTree structure
  getBranchesStage1(inTree,&myTree,mcData);
  
  //Find out how many events are in the tree and set to nToRun
  if(nToRun<0) nToRun = inTree->GetEntries();

  //Print the number of events to process in the terminal
  cout<<"Number of events to process = "<<nToRun<<endl;

  double mPi = 0.139568;
  double mPi0 = 0.13497;
  double mPro = 0.9381;
  double mK = 0.493677;
  double mK0 = 0.4977;
  double mNeut = 0.9465;
  
  //cout<<"test:1"<<endl;
  for(Int_t i=0;i<nToRun;i++){
    //cout<<"test:2"<<endl;
    inTree->GetEntry(i);
    //Print every 1000 events                                          
    if (i%1000 == 0) cout<<"Events processed = "<<i<<endl;
    //Find out how many particles are in the event                      
    //cout<<"test:3"<<endl;
    int nREC_Particle = myTree.REC_Particle_pid->size();
    int nRECFT_Particle = myTree.RECFT_Particle_pid->size();
    int nREC_Calorimeter = myTree.REC_Calorimeter_pindex->size();
    //Find out how many cherenkov entries per event
    //int nCher = myTree.REC_Cherenkov_index->size();
    //cout<<"test:4"<<endl;
  //Setting up Bucket of Information 
    myBOI_t boi;
    vector <myBOI_t> proBOI;
    vector <myBOI_t> phoBOI;
    vector <myBOI_t> pipBOI;
    vector <myBOI_t> pimBOI;
    vector <myBOI_t> neutBOI;
    vector <myBOI_t> elecBOI;
    vector <myBOI_t> kpBOI;
    vector <myBOI_t> kmBOI;
    vector <myBOI_t> phoAllBOI;

    //Start MC stuff
    vector <myBOI_t> proBOImc;
    vector <myBOI_t> kpBOImc;
    vector <myBOI_t> pimBOImc;
    vector <myBOI_t> elecBOImc;

    int nMC_Event = 0;    
    if (mcData == true) nMC_Event = myTree.MC_Event_ebeam->size();    
    double eBeamMC = -1;
    if (nMC_Event == 1) eBeamMC = myTree.MC_Event_ebeam->at(0);
    int nMC_Particle = 0;
    if (mcData == true) nMC_Particle = myTree.MC_Particle_pid->size();    

    //cout<<"nMC_Particle = "<<nMC_Particle<<endl;

    for(int index = 0; index<nMC_Particle; index++){
      int pid   = myTree.MC_Particle_pid->at(index);
      //cout<<"MC pid = "<<pid<<endl;
      double px   = myTree.MC_Particle_px->at(index);
      double py   = myTree.MC_Particle_py->at(index);
      double pz   = myTree.MC_Particle_pz->at(index);
      double p   = sqrt(pow(px,2)+pow(py,2)+pow(pz,2));
      double eval = 0;
      if (pid == 2212) {
	eval = sqrt(pow(p,2)+pow(mPro,2));
        boi.mo.SetPxPyPzE(px,py,pz,eval);
	proBOImc.push_back(boi);
      }
      if (pid == 321) {
	eval = sqrt(pow(p,2)+pow(mK,2));
        boi.mo.SetPxPyPzE(px,py,pz,eval);
	kpBOImc.push_back(boi);
      }
      if (pid == -211) {
	eval = sqrt(pow(p,2)+pow(mPi,2));
        boi.mo.SetPxPyPzE(px,py,pz,eval);
	pimBOImc.push_back(boi);
      }
      if (pid == 11) {
	eval = sqrt(pow(p,2)+pow(0.000511,2));
        boi.mo.SetPxPyPzE(px,py,pz,eval);
	elecBOImc.push_back(boi);
      }

    }
    int nproMC = proBOImc.size();
    int nkpMC = kpBOImc.size();
    int npimMC = pimBOImc.size();
    int nelecMC = elecBOImc.size();

    int nMC = nkpMC + nproMC + npimMC;

    for(int index = 0; index<nRECFT_Particle; index++){
      //cout<<"test:6"<<endl;
      int pid   = myTree.RECFT_Particle_pid->at(index);
      int partStat = myTree.RECFT_Particle_status->at(index);
      int partStatDet = partStat - partStat%1000;
      boi.sciDet = partStatDet;
      
      boi.nphe = 0;

      //See https://clasweb.jlab.org/wiki/index.php/CLAS12_DSTs

      double px = myTree.REC_Particle_px->at(index);
      double py = myTree.REC_Particle_py->at(index);
      double pz = myTree.REC_Particle_pz->at(index);
      double p  = sqrt(pow(px,2)+pow(py,2)+pow(pz,2));
      double beta   = myTree.RECFT_Particle_beta->at(index);
      double gamma = -1;
      if (beta != 1)  gamma = sqrt(1/(1-beta*beta));
      double massSQ = pow(p/(gamma*beta),2);
      double mass = -1;
      double eval = -1;
      double nphe = boi.nphe;
      
      boi.beta = beta;
      boi.p = p;
         
      if (massSQ>0) mass = sqrt(massSQ);
      boi.mass = mass;
      //eft -> electron seen in forward tagger
      if(pid == 11){

	double eft = sqrt(pow(p,2)+pow(0.000511,2));
     	boi.mo.SetPxPyPzE(px,py,pz,p);
	double eftNew = eft-0.03689+0.1412*eft-0.04316*pow(eft,2)+0.007046*pow(eft,3)-0.0004055*pow(eft,4);
	double eftFrac = eftNew/eft;
	eftFrac = 1.0888 - 0.0189 * eft + 0.0031 * pow(eft,2)- 0.0002 * pow(eft,3);
	TLorentzVector p4VecCor(px*eftFrac,py*eftFrac,pz*eftFrac,p*eftFrac);
	boi.mo = p4VecCor;
	elecBOI.push_back(boi);
      }
            
      if (pid == 2212 && mass>0) {
	hBetaVsPPro->Fill(p,beta);
	eval = sqrt(pow(p,2)+pow(mPro,2));
	boi.mo.SetPxPyPzE(px,py,pz,eval);
	proBOI.push_back(boi);
      }
      hBetaVsP->Fill(p,beta);

      //pi+
      if (pid == 211 && mass>0) {
        eval = sqrt(pow(p,2)+pow(mPi,2));
        boi.mo.SetPxPyPzE(px,py,pz,eval);
	pipBOI.push_back(boi);
      }
      //pi-
      //cout<<"pid = "<<pid<<endl;
      if (pid == -211 && mass>0) {
	eval = sqrt(pow(p,2)+pow(mPi,2));
        boi.mo.SetPxPyPzE(px,py,pz,eval);
	pimBOI.push_back(boi);
      }
      //neut
      if (pid == 2112 && mass>0) {
        eval = sqrt(pow(p,2)+pow(mNeut,2));
        boi.mo.SetPxPyPzE(px,py,pz,eval);
	neutBOI.push_back(boi);
      }
      //K+
      if (pid == 321 && mass>0) {
	if(beta>0.855) {
	  eval = sqrt(pow(p,2)+pow(mK,2));
	  boi.mo.SetPxPyPzE(px,py,pz,eval);
	  if (partStatDet == 2000){
	    kpBOI.push_back(boi);
	  }
	}
      }
      //K-
      if (pid == -321 && mass>0 && beta>0.855) {
	eval = sqrt(pow(p,2)+pow(mK,2));
	boi.mo.SetPxPyPzE(px,py,pz,eval);
	if (partStatDet == 2000) kmBOI.push_back(boi);
      }
    }

    int npro = proBOI.size();
    int npho = phoBOI.size();
    int npip = pipBOI.size();
    int npim = pimBOI.size();
    int nneut = neutBOI.size();
    int nkp = kpBOI.size();
    int nelec = elecBOI.size();
    int nkm = kmBOI.size();

    int nPart = nkp + npro + npim;


    //Here is the reaction stuff
    if (npro == 1 && npim == 1 && nelec == 1){
      TLorentzVector mopropim = proBOI[0].mo + pimBOI[0].mo;
      double masspropim = mopropim.Mag();
      hmasspropim->Fill(masspropim);
      if (nkp == 1){
	hmasspropim2->Fill(masspropim);
      }
    }

  }
  
  outFile->cd();

  hBetaVsP->Write();
  hBetaVsPPro->Write();

  hmasspropim->Write();
  hmasspropim2->Write();
  outFile->Write();

  return 0;
}
void printUsage(){
  fprintf(stderr,"\nUsage:");
  fprintf(stderr,"\nstage2 [-switches] <inputFile>\n");
  fprintf(stderr,"\nSWITCHES:\n");
  fprintf(stderr,"-h\tPrint this message\n");
  fprintf(stderr,"-m\tNeed for Monte Carlo Events\n");
  fprintf(stderr,"-o<arg>\tOutFileName. Default is myOutFile.root\n\n");

  cout<<"The current default operation is equivalent to the command:"<<endl;
  cout<<"myTreeCodeV2 -omyOutFile.root inputFile ,"<<endl;
  cout<<"where \"inputFile\" must be supplied by the user\n"<<endl;

  exit(0);
}