#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 *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("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);

  //DEFINE OUTPUT TREE                                                          
  TTree outTree("mass_pi0kpkm__B4_M7_Tree","mass_pi0kpkm__B4_M7_Tree");

  //DEFINE OUT BRANCHES                               
  setBranchesHitToGo(&outTree,&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;


    if (beam_energy < 5.2) continue;
    //if (beam_energy > 8.8) continue;

    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 beam4Vec(0.,0.,beam_energy,beam_energy);
    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);

    angleBOI_t angleBOI;

    angleBOI = getBoosted(kp4Vec,km4Vec,pi4Vec,beam4Vec);
    
    myTree.cosTheta     = angleBOI.cosTheta;
    myTree.cosTheta_H   = angleBOI.cosTheta_H;
    myTree.cosThetaGJ   = angleBOI.cosThetaGJ;
    myTree.cosThetaGJ_H = angleBOI.cosThetaGJ_H;

    myTree.theta     = angleBOI.theta;
    myTree.theta_H   = angleBOI.theta_H;
    myTree.thetaGJ   = angleBOI.thetaGJ;
    myTree.thetaGJ_H = angleBOI.thetaGJ_H;

    myTree.phi     = angleBOI.phi;
    myTree.phi_H   = angleBOI.phi_H;
    myTree.phiGJ   = angleBOI.phiGJ;
    myTree.phiGJ_H = angleBOI.phiGJ_H;
    


    outTree.Fill();
  }


  outFile->cd();
  hMassKpKm->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);
}

angleBOI_t getBoosted(TLorentzVector p1,TLorentzVector p2,TLorentzVector p3,TLorentzVector beam){
  
  TLorentzVector resonance = p1 + p2 + p3;
  TLorentzVector daughter1 = p1 + p2;
  TLorentzVector daughter2 = p1;
  TLorentzVector daughter3 = p2;
  
  TLorentzVector target(0,0,0,0.938);
  TLorentzVector cm = beam+target;
  
  // beam polarization vector
  //TVector3 eps(cos(polAngle*TMath::DegToRad()), sin(polAngle*TMath::DegToRad()), 0.0); 
  //TLorentzVector eps4Vec(cos(polAngle*TMath::DegToRad()), sin(polAngle*TMath::DegToRad()),0.0,0.0);
  
  //Boost to CM frame
  TVector3 cmRestBoost = cm.BoostVector();
  TLorentzVector beam_cm = beam;
  TLorentzVector resonance_cm = resonance;
  TLorentzVector daughter1_cm = daughter1;
  TLorentzVector daughter2_cm = daughter2;
  TLorentzVector daughter3_cm = daughter3;
  beam_cm.Boost(-1.0*cmRestBoost);
  resonance_cm.Boost(-1.0*cmRestBoost);
  daughter1_cm.Boost(-1.0*cmRestBoost);
  daughter2_cm.Boost(-1.0*cmRestBoost);
  daughter3_cm.Boost(-1.0*cmRestBoost);

  //Boost to resonance frame
  TVector3 resRestBoost = resonance_cm.BoostVector();
  TLorentzVector beam_res = beam_cm;
  TLorentzVector resonance_res = resonance_cm;
  TLorentzVector daughter1_res = daughter1_cm;
  TLorentzVector daughter2_res = daughter2_cm;
  TLorentzVector daughter3_res = daughter3_cm;
  beam_res.Boost(-1.0*resRestBoost);
  resonance_res.Boost(-1.0*resRestBoost);
  daughter1_res.Boost(-1.0*resRestBoost);
  daughter2_res.Boost(-1.0*resRestBoost);
  daughter3_res.Boost(-1.0*resRestBoost);

  //Boost to daughter frame
  TVector3 vRestBoost = daughter1_res.BoostVector();
  TLorentzVector daughter2_v = daughter2_res;
  TLorentzVector daughter3_v = daughter3_res;
  daughter2_v.Boost(-1.0*vRestBoost);
  daughter3_v.Boost(-1.0*vRestBoost);

  TVector3 beam_cm_vec = beam_cm.Vect().Unit();
  TVector3 daughter1_res_vec = daughter1_res.Vect().Unit();
  TVector3 daughter2_vec = daughter2_v.Vect().Unit();
  TVector3 daughter3_vec = daughter3_v.Vect().Unit();

  //Define coordinate system in res frame using cm-frame variables
  //z-axis is in direction of resonance_cm. Production plane
  //of beam/res contains z and x
  TVector3 z = resonance_cm.Vect().Unit();
  TVector3 y = (beam_cm_vec.Cross(z)).Unit();
  TVector3 x = (y.Cross(z)).Unit();

  TVector3 Angles(daughter1_res_vec.Dot(x),
		  daughter1_res_vec.Dot(y),
		  daughter1_res_vec.Dot(z));

  double cosTheta = Angles.CosTheta();
  double theta = Angles.Theta();
  double phi = Angles.Phi();

  //H rotation
  TVector3 z_H = daughter1_res.Vect().Unit();
  TVector3 y_H = (z.Cross(z_H)).Unit();
  TVector3 x_H = (y_H.Cross(z_H)).Unit();


  TVector3 Angles_H(daughter2_vec.Dot(x_H),
                    daughter2_vec.Dot(y_H),
                    daughter2_vec.Dot(z_H));


  double cosTheta_H = Angles_H.CosTheta();
  double theta_H = Angles_H.Theta();
  double phi_H = Angles_H.Phi();

  //////////////////////////////////////////////////////////////
  //Rest-frame of the meson-resonance: Gottfried-Jackson frame//
  //////////////////////////////////////////////////////////////

  //TVector3 xLabUnit(1,0,1);
  //TVector3 yLabUnit(0,1,0);
  //TVector3 zLabUnit(0,0,1);
  TVector3 zGJ = beam_res.Vect().Unit();
  TVector3 yGJ = y;
  TVector3 xGJ = (yGJ.Cross(zGJ)).Unit();

  TVector3 AnglesGJ(daughter1_res_vec.Dot(xGJ),
		    daughter1_res_vec.Dot(yGJ),
		    daughter1_res_vec.Dot(zGJ));

  double cosThetaGJ = AnglesGJ.CosTheta();
  double thetaGJ = AnglesGJ.Theta();
  double phiGJ = AnglesGJ.Phi();

  //H rotation
  TVector3 zGJ_H = daughter1_res.Vect().Unit();
  TVector3 yGJ_H = (zGJ.Cross(zGJ_H)).Unit();
  TVector3 xGJ_H = (yGJ_H.Cross(zGJ_H)).Unit();

  TVector3 AnglesGJ_H(daughter2_vec.Dot(xGJ_H),
		    daughter2_vec.Dot(yGJ_H),
		    daughter2_vec.Dot(zGJ_H));

  double cosThetaGJ_H = Angles_H.CosTheta();
  double thetaGJ_H = AnglesGJ_H.Theta();
  double phiGJ_H = AnglesGJ_H.Phi();

  angleBOI_t angleBOI;
  angleBOI.theta = theta;
  angleBOI.cosTheta = cosTheta;
  angleBOI.phi = phi;
  angleBOI.theta_H = theta_H;
  angleBOI.cosTheta_H = cosTheta_H;
  angleBOI.phi_H = phi_H;
  angleBOI.thetaGJ = thetaGJ;
  angleBOI.cosThetaGJ = cosThetaGJ;
  angleBOI.phiGJ = phiGJ;
  angleBOI.thetaGJ_H = thetaGJ_H;
  angleBOI.cosThetaGJ_H = cosThetaGJ_H;
  angleBOI.phiGJ_H = phiGJ_H;

  return angleBOI;

}