#include <string.h>
#include <iostream>
#include <vector>
#include "TTree.h"
#include "TFile.h"
#include "TLorentzVector.h"
#include "TVector3.h"
#include "TH1.h"
#include "TH2.h"
#include "TF1.h"
#include "myTree.h"
#include <bitset>
#include "TDatabasePDG.h"
#include "TParticlePDG.h"
 
 
//FUNCTION PROTOTYPES
void printUsage(); // print usage

//MAIN
int main(int argc, char **argv) {
  
  char  inFileName[120];
  char  outFileName[120];
  char *argptr;
  int nToRun = -1;
  //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 'o':
        strcpy(outFileName,++argptr);
        break;
      default:
        printUsage();
        break;
      }
    } else {
      strcpy(inFileName,argptr);
    }
  }

  TH1D *hMassKpKm     = new TH1D("hMassKpKm","",1000,0.8,1.8);


  //Setup the input tree
  ////Get file object
  TFile* inFile=new TFile(inFileName);
  ////Get tree out of the file
  TTree *inTree = (TTree*)inFile->Get("clas12");

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

  //Define the tree structure
  myTree_t myTree;
  //Get the tree branches into my tree structure
  getBranchesStage1(inTree,&myTree);
  
  if (nToRun< 0) nToRun = inTree->GetEntries();
  cout<<"Number of events to process = "<<nToRun<<endl;
  
  double mPro  = 0.938;
  double mPip  = 0.13957;
  double mPim  = 0.13957;
  double mKp   = 0.493677;
  double mLambda  = 1.115;
  double mElec = 0.000511;
  
  for(Int_t i=0;i<nToRun;i++){
    inTree->GetEntry(i);
    
    if (i%50000 == 0) cout<<"Events processed = "<<i<<endl;
    
    vector<myBucketOfInformation_t> proBOI;
    vector<myBucketOfInformation_t> pipBOI;
    vector<myBucketOfInformation_t> pimBOI;
    vector<myBucketOfInformation_t> kpBOI;
    vector<myBucketOfInformation_t> kmBOI;
    vector<myBucketOfInformation_t> elecBOI;
    vector<myBucketOfInformation_t> phoBOI;
    
    vector<myBucketOfInformation_t> proFTBOI;
    vector<myBucketOfInformation_t> pipFTBOI;
    vector<myBucketOfInformation_t> pimFTBOI;
    vector<myBucketOfInformation_t> kpFTBOI;
    vector<myBucketOfInformation_t> kmFTBOI;
    
    myBucketOfInformation_t tmpBOI;
   
    int nREC_Particle = myTree.RECFT_Particle_pid->size();
    int nRECFT_Particle = myTree.RECFT_Particle_pid->size();
    
    for(int index = 0; index<nREC_Particle; index++){

      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 vx = myTree.REC_Particle_vx->at(index);
      double vy = myTree.REC_Particle_vy->at(index);
      double vz = myTree.REC_Particle_vz->at(index);
      double vt = myTree.RECFT_Particle_vt->at(index);

      TLorentzVector p4Vec(myTree.REC_Particle_px->at(index),myTree.REC_Particle_py->at(index),
			    myTree.REC_Particle_pz->at(index),p);

      double beta   = myTree.REC_Particle_beta->at(index);
      double betaFT = myTree.RECFT_Particle_beta->at(index);
      double mass = p * sqrt((1 / pow(beta,2)) - 1);
      double chi2 = myTree.RECFT_Particle_chi2pid->at(index);
      double chi2FT = myTree.REC_Particle_chi2pid->at(index);

      int pid = myTree.REC_Particle_pid->at(index);
      int pidFT = myTree.RECFT_Particle_pid->at(index);
      
      int partStat = myTree.RECFT_Particle_status->at(index);
      int partStatDet = partStat - partStat%1000;
      tmpBOI.sciDet = partStatDet;

      double pimCor = mPip/(0.140 + (chi2 * 0.0125));
      double pipCor = mPip/(0.140 + (chi2 * 0.034));
      double kpCor = mKp/(mKp + (chi2 * 0.034));
      double proCor = mPro/(0.938 + chi2*0.034);
      
      //Proton 2212
      if (pid == 2212) {
	double eTmp = sqrt(pow(mPro,2)+ pow(p,2));
	p4Vec.SetE(eTmp);
	tmpBOI.momentum = p4Vec;
	tmpBOI.momentumCor = p4Vec;
	if (partStatDet == 4000){
	  double eTmpCor = sqrt(pow(mPro,2)+ pow(p*proCor,2));
	  double massCor = p*proCor * sqrt((1 / pow(beta,2)) - 1);
	  TLorentzVector p4VecCor(px*proCor,py*proCor,pz*proCor,eTmpCor);
	  tmpBOI.momentumCor = p4VecCor;
	}
        proBOI.push_back(tmpBOI);
      }
      //PiPlus 211
      if (pid == 211) {
	double eTmp = sqrt(pow(mPip,2)+ pow(p,2));
	p4Vec.SetE(eTmp);
	tmpBOI.momentum = p4Vec;
	tmpBOI.momentumCor = p4Vec;
	if (partStatDet == 4000){
	  double eTmpCor = sqrt(pow(mPip,2)+ pow(p*pipCor,2));
	  double massCor = p*pimCor * sqrt((1 / pow(beta,2)) - 1);
	  TLorentzVector p4VecCor(px*pipCor,py*pipCor,pz*pipCor,eTmpCor);
	  tmpBOI.momentumCor = p4VecCor;
	}
        pipBOI.push_back(tmpBOI);
      }
      //PiMinus  -211
      if (pid == -211) {
	double eTmp = sqrt(pow(mPip,2)+ pow(p,2));
	p4Vec.SetE(eTmp);
	tmpBOI.momentum = p4Vec;
	tmpBOI.momentumCor = p4Vec;
	if (partStatDet == 4000){
	  double eTmpCor = sqrt(pow(mPim,2)+ pow(p*pimCor,2));
	  double massCor = p*pimCor * sqrt((1 / pow(beta,2)) - 1);
	  TLorentzVector p4VecCor(px*pimCor,py*pimCor,pz*pimCor,eTmpCor);
	  tmpBOI.momentumCor = p4VecCor;
	}
        pimBOI.push_back(tmpBOI);
      }
      //KPlus 321
      if (pid == 321) {
	double eTmp = sqrt(pow(mKp,2)+ pow(p,2));
	p4Vec.SetE(eTmp);
	tmpBOI.momentum = p4Vec;
        if (partStatDet == 2000) kpBOI.push_back(tmpBOI);
      }
      //KMinus -321
      if (pid == -321) {
	double eTmp = sqrt(pow(mKp,2)+ pow(p,2));
	p4Vec.SetE(eTmp);
	tmpBOI.momentum = p4Vec;
        if (partStatDet == 2000) kmBOI.push_back(tmpBOI);
      }
      //Electron -11
      if (pid == -11) {
	tmpBOI.momentum = p4Vec;
	double eft = p;
	double eftFrac = (eft - 0.0004*pow(eft,4) + 0.0071*pow(eft,3) - 0.0432*pow(eft,2) + 0.1356*eft - 0.0257)/eft;

	TLorentzVector p4VecCor(px*eftFrac,py*eftFrac,pz*eftFrac,p*eftFrac);
	tmpBOI.momentumCor = p4VecCor;
        elecBOI.push_back(tmpBOI);
      }
    }

    int nPro  = proBOI.size();
    int nPip  = pipBOI.size();
    int nPim  = pimBOI.size();
    int nKp   = kpBOI.size();
    int nKm   = kmBOI.size();
    int nElec = elecBOI.size();
    int nPhoBOI = phoBOI.size();
    
    
    double eBeamVal = 10.6041;
    TLorentzVector eBeam(0.0,0.0,eBeamVal,eBeamVal);
    TLorentzVector target(0.0,0.0,0.0,mPro);

    if (nElec == 1 && nKp ==1 && nKm == 1){
	TLorentzVector moKpKm = kpBOI[0].momentum + kmBOI[0].momentum;
	hMassKpKm->Fill(moKpKm.Mag());

    } 
  }


  outFile->cd();
  hMassKpKm->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,"-o<arg>\tOutFileName. Default is myOutFile.root\n\n");

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

  exit(0);
}