#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"
 
struct angleBOI_t {
  Float_t theta;
  Float_t cosTheta;
  Float_t phi;
  Float_t theta_H;
  Float_t cosTheta_H;
  Float_t phi_H;
  Float_t thetaGJ;
  Float_t cosThetaGJ;
  Float_t phiGJ;
  Float_t thetaGJ_H;
  Float_t cosThetaGJ_H;
  Float_t phiGJ_H;
}; 

//FUNCTION PROTOTYPES
void printUsage(); // print usage
angleBOI_t getBoosted(TLorentzVector p1,TLorentzVector p2,TLorentzVector p3,TLorentzVector beam);

//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 *hMassKpKmPi0     = new TH1D("hMassKpKmPi0","",100,1.1,1.35);
  TH1D *hMassKpKm     = new TH1D("hMassKpKm","",100,0.98,1.1);
  TH1D *hMassPi0     = new TH1D("hMassPi0","",100,0.11,0.16);
  TH1D *hMassKpStar     = new TH1D("hMassKpStar","",100,0.6,0.9);
  TH1D *hMassKmStar     = new TH1D("hMassKmStar","",100,0.6,0.9);
  TH1D *hMMSq     = new TH1D("hMMSq","",100,-0.000001,0.000001);
  TH1D *hMM     = new TH1D("hMM","",100,-0.0,0.002);

  TH1D *hMoKp     = new TH1D("hMoKp","",100,0.0,4.0);
  TH1D *hMoKm     = new TH1D("hMoKm","",100,0.0,4.0);

  TH1D *hEGamma     = new TH1D("hEGamma","",100,7.0,10.0);

  TH1D *hCosGJ     = new TH1D("hCosGJ","",10,-1.0,1.0);
  TH1D *hPhiGJ     = new TH1D("hPhiGJ","",18,0.0,360.0);
  TH1D *hCosH     = new TH1D("hCosH","",10,-1.0,1.0);
  TH1D *hPhiH     = new TH1D("hPhiH","",18,0.0,360.0);


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

  //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;
  
  
  for(Int_t i=0;i<nToRun;i++){
    if (i%50000 == 0) cout<<"Events processed = "<<i<<endl;

    inTree->GetEntry(i);


    Float_t beam_energy = myTree.beam_energy;

    Float_t p_px = myTree.p_px;
    Float_t p_py = myTree.p_py;
    Float_t p_pz = myTree.p_pz;
    Float_t p_E  = myTree.p_E;

    Float_t kp_px = myTree.kp_px;
    Float_t kp_py = myTree.kp_py;
    Float_t kp_pz = myTree.kp_pz;
    Float_t kp_E  = myTree.kp_E;

    Float_t km_px = myTree.km_px;
    Float_t km_py = myTree.km_py;
    Float_t km_pz = myTree.km_pz;
    Float_t km_E  = myTree.km_E;

    Float_t pi0_px = myTree.pi0_px;
    Float_t pi0_py = myTree.pi0_py;
    Float_t pi0_pz = myTree.pi0_pz;
    Float_t pi0_E  = myTree.pi0_E;

    TLorentzVector target4Vec(0.0,0.0,0.0,0.938);
    TLorentzVector beam4Vec(0.0,0.0,beam_energy,beam_energy);
    TLorentzVector p4Vec(p_px,p_py,p_pz,p_E);

    TLorentzVector kp4Vec(kp_px,kp_py,kp_pz,kp_E);
    TLorentzVector km4Vec(km_px,km_py,km_pz,km_E);
    TLorentzVector pi4Vec(pi0_px,pi0_py,pi0_pz,pi0_E);
    TLorentzVector missing4Vec = target4Vec + beam4Vec - p4Vec - kp4Vec - km4Vec - pi4Vec;

    TLorentzVector kpStar4Vec = kp4Vec + pi4Vec;
    TLorentzVector kmStar4Vec = km4Vec + pi4Vec;
    
    double phiGJdeg = myTree.phiGJ*180/3.1419;
    if (phiGJdeg < 0) phiGJdeg += 360.0;

    double phiGJ_Hdeg = myTree.phiGJ_H*180/3.1419;
    if (phiGJ_Hdeg < 0) phiGJ_Hdeg += 360.0;
    
    hMassKpKmPi0->Fill(myTree.m_kpkmpi0);
    hMassKpKm->Fill(myTree.m_kpkm);
    hMassPi0->Fill(pi4Vec.Mag());
    hMassKpStar->Fill(kpStar4Vec.Mag());
    hMassKmStar->Fill(kmStar4Vec.Mag());
    hMM->Fill(missing4Vec.Mag());
    hMMSq->Fill(missing4Vec.Mag2());

    hEGamma->Fill(beam_energy);
    hMoKp->Fill(kp4Vec.P());
    hMoKm->Fill(km4Vec.P());

    hCosGJ->Fill(myTree.cosThetaGJ);
    hPhiGJ->Fill(phiGJdeg);
    hCosH->Fill(myTree.cosThetaGJ_H);
    hPhiH->Fill(phiGJ_Hdeg);

    if (kp4Vec.P()>3.0) continue;
    if (km4Vec.P()>3.0) continue;
    if (pi4Vec.Mag()<0.12) continue;
    if (pi4Vec.Mag()>0.15) continue;
    if (myTree.m_kpkmpi0 < 1.233701) continue;
    if (myTree.m_kpkmpi0 > 1.330099) continue;
    if (kpStar4Vec.Mag() > 0.82618729) continue;
    if (kmStar4Vec.Mag() > 0.82618729) continue;
    if (beam_energy < 8.2) continue;
    if (beam_energy > 8.8) continue;

    //Both kaons need to have momentum < 3.0
    //0.12 < Mass[γγ] < 0.15
    //1.233701 < Mass[KKπ] < 1.330099 (these numbers might seem random, 
    //but I got them using the PDG center and width of the f1(1285) and 
    //thinking in terms of standard deviations)
    //Mass[Kπ] < 0.82618729 (for both kaons)
    //Mass[KK] < 1.0104391 
    //Missing mass < 0.0015
    //8.2 < Beam energy < 8.8 (newest one)

  }


  outFile->cd();
  hMassKpKm->SetMinimum(0);
  hMassKpKmPi0->SetMinimum(0);
  hMassPi0->SetMinimum(0);
  hMassKpStar->SetMinimum(0);
  hMassKmStar->SetMinimum(0);

  hMoKp->SetMinimum(0);
  hMoKm->SetMinimum(0);
  hEGamma->SetMinimum(0);

  hCosGJ->SetMinimum(0);
  hPhiGJ->SetMinimum(0);
  hCosH->SetMinimum(0);
  hPhiH->SetMinimum(0);

  hMassKpKm->Write();
  hMassKpKmPi0->Write();
  hMassPi0->Write();
  hMassKpStar->Write();
  hMassKmStar->Write();
  hMM->Write();
  hMMSq->Write();

  hMoKp->Write();
  hMoKm->Write();
  hEGamma->Write();

  hCosGJ->Write();
  hPhiGJ->Write();
  hCosH->Write();
  hPhiH->Write();
  //outTree.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);
}