//this does the selection with TCuts

using namespace std;

#include "TFile.h"
#include "TTree.h"
#include "TCut.h"
#include "TEventList.h"
#include <iostream>
#include "TMath.h"
#include <TChain.h>
#include <string.h>
#include <TROOT.h>
#include <TEntryList.h>
#include <TLorentzVector.h>
#include <TVector3.h>

#define Kaon_M_Official 493.667

//calculates the cos helicity of particle
Float_t coshel(TLorentzVector particle, TLorentzVector parent, TLorentzVector grandparent) {
  // The helicity of a particle - Re: http://slac.stanford.edu/BFROOT/www/doc/workbook/analysis/analysis.html
  TVector3 boosttoparent = -(parent.BoostVector());

  particle.Boost(boosttoparent);
  grandparent.Boost(boosttoparent);

  TVector3 particle3 = particle.Vect();
  TVector3 grandparent3 = grandparent.Vect();

  Float_t numerator = particle3.Dot(grandparent3);
  Float_t denominator = (particle3.Mag())*(grandparent3.Mag());
  Float_t temp = numerator/denominator;

  return temp;
}

Float_t cospointing(TVector3 particle, TVector3 partDirection) {
  // Cosine of the angle between particle momentum and its direction
  // direction : eg. vector between particle and its previous vertex
  
  Float_t numerator = particle.Dot(partDirection);
  Float_t denominator = (particle.Mag())*(partDirection.Mag());
  Float_t temp = numerator/denominator;

  return temp;
}

//adding some branches to the tree to be cut on
void coshel_and_Bdcosptng(TTree *tree)
{
  //vars to be added
  Float_t Kstar_coshel;
  Float_t Bdcosptng;
  
  //vars needed
  Float_t KstarK_REFITTED_PX;
  Float_t KstarK_REFITTED_PY;
  Float_t KstarK_REFITTED_PZ;
  Float_t KstarK_e;
  
  Float_t Kstar_PX;
  Float_t Kstar_PY;
  Float_t Kstar_PZ;
  Float_t Kstar_PE;
  
  Float_t Bd_PX;
  Float_t Bd_PY;
  Float_t Bd_PZ;
  Float_t Bd_PE;

  Float_t Bd_ENDVERTEX_X;
  Float_t Bd_ENDVERTEX_Y;
  Float_t Bd_ENDVERTEX_Z;
  Float_t Bd_OWNPV_X;
  Float_t Bd_OWNPV_Y;
  Float_t Bd_OWNPV_Z;

  TLorentzVector kFromKStar(0,0,0,0);
  TLorentzVector KStarParent(0,0,0,0);
  TLorentzVector BGranPa(0,0,0,0);

  TVector3 BdMomentum(0,0,0);
  TVector3 BdDirection(0,0,0);  
  
  TBranch *Kstar_coshelBranch = tree->Branch("Kstar_coshel", &Kstar_coshel, "Kstar_coshel/F");
  TBranch *BdcosptngBranch = tree->Branch("Bdcosptng", &Bdcosptng, "Bdcosptng/F");
  
  tree->SetBranchAddress("KstarK_REFITTED_PX", &KstarK_REFITTED_PX);
  tree->SetBranchAddress("KstarK_REFITTED_PY", &KstarK_REFITTED_PY);
  tree->SetBranchAddress("KstarK_REFITTED_PZ", &KstarK_REFITTED_PZ);

  tree->SetBranchAddress("Kstar_PX", &Kstar_PX);
  tree->SetBranchAddress("Kstar_PY", &Kstar_PY);  
  tree->SetBranchAddress("Kstar_PZ", &Kstar_PZ);  
  tree->SetBranchAddress("Kstar_PE", &Kstar_PE);  

  tree->SetBranchAddress("Bd_PX", &Bd_PX);
  tree->SetBranchAddress("Bd_PY", &Bd_PY);  
  tree->SetBranchAddress("Bd_PZ", &Bd_PZ);  
  tree->SetBranchAddress("Bd_PE", &Bd_PE);  

  tree->SetBranchAddress("Bd_ENDVERTEX_X", &Bd_ENDVERTEX_X);
  tree->SetBranchAddress("Bd_ENDVERTEX_Y", &Bd_ENDVERTEX_Y);
  tree->SetBranchAddress("Bd_ENDVERTEX_Z", &Bd_ENDVERTEX_Z);
  tree->SetBranchAddress("Bd_OWNPV_X", &Bd_OWNPV_X);
  tree->SetBranchAddress("Bd_OWNPV_Y", &Bd_OWNPV_Y);
  tree->SetBranchAddress("Bd_OWNPV_Z", &Bd_OWNPV_Z);
  
  //read the number of entries in the tree
  Long64_t nentries = tree->GetEntries();  
  std::cout << "Num Entries: " << nentries << std::endl;
  
  for (Int_t i = 0; i < nentries; i++){
    tree->GetEntry(i);
    
    KstarK_e = sqrt(Kaon_M_Official*Kaon_M_Official + KstarK_REFITTED_PX*KstarK_REFITTED_PX
                    + KstarK_REFITTED_PY*KstarK_REFITTED_PY + KstarK_REFITTED_PZ *KstarK_REFITTED_PZ);
    kFromKStar = TLorentzVector(KstarK_REFITTED_PX, KstarK_REFITTED_PY, KstarK_REFITTED_PZ, KstarK_e);
    KStarParent = TLorentzVector(Kstar_PX, Kstar_PY, Kstar_PZ, Kstar_PE);
    BGranPa = TLorentzVector(Bd_PX, Bd_PY, Bd_PZ, Bd_PE);
    Kstar_coshel = coshel(kFromKStar, KStarParent, BGranPa);
    Kstar_coshelBranch->Fill();
    
    BdMomentum = TVector3(Bd_PX, Bd_PY, Bd_PZ);
    BdDirection = TVector3(Bd_ENDVERTEX_X-Bd_OWNPV_X, Bd_ENDVERTEX_Y-Bd_OWNPV_Y, Bd_ENDVERTEX_Z-Bd_OWNPV_Z);
    Bdcosptng = cospointing(BdMomentum, BdDirection);
    BdcosptngBranch->Fill();
    
  }
  
  
}

int main() 
{
  cout << "Hello World! I am going to perform the selection on Bd -> DK* !!" << endl;
  cout << "I'm using TCuts because it's easier." << endl;
  gROOT->Time();
  TChain *fChain; 
  fChain = new TChain("B02D0Kstar_D02KPi_Tuple/DecayTree");
  string type = "MC";
  
  // Actually create the chain
  if ( type == "Data" ) {
    fChain->Add("/data/lhcb/users/smithe/DataNtuples/Data_2011/DVNtuples_Reco10_Stripping13b_MagUp_1_11June2011.root");
    fChain->Add("/data/lhcb/users/smithe/DataNtuples/Data_2011/DVNtuples_Reco10_Stripping13b_MagUp_4_30June2011.root");
    fChain->Add("/data/lhcb/users/smithe/DataNtuples/Data_2011/DVNtuples_Reco10_Stripping13b_MagUp_5_08Jul72011.root");
    fChain->Add("/data/lhcb/users/smithe/DataNtuples/Data_2011/DVNtuples_Reco10_Stripping13b_MagUp_5June2011.root");
    fChain->Add("/data/lhcb/users/smithe/DataNtuples/Data_2011/DVNtuples_Reco10_Stripping13b_MagDown_4_30June2011.root");
    fChain->Add("/data/lhcb/users/smithe/DataNtuples/Data_2011/DVNtuples_Reco10_Stripping13b_MagDown_5_08Jul72011.root");
    //fChain->Add("/data/lhcb/users/smithe/DataNtuples/Data_2011/");
  } else if ( type == "MC" ) {
    fChain->Add("/data/lhcb/users/smithe/MC10/MC10Ntuples_Bs2D0KstarAll-Stripping12b.root");
  }

  //mass window cuts
  TCut mass_D0 = "abs(D0_MM-1864.8)<20";
  TCut mass_Kstar = "abs(Kstar_MM-775.49)<50";

  
  //PIDK cuts
  TCut D0Kpid = "D0K_PIDK>0.0";
  TCut D0Pipid = "(D0Pi_PIDK<4.0)&&(D0Pi_PIDK>-999.0)";
  TCut KstarPipid = "(KstarPi_PIDK<3.0)&&(KstarPi_PIDK>-999.0)";
  TCut KstarKpid = "(KstarK_PIDK>3.0)";
  
  TCut pidcut = D0Kpid && D0Pipid && KstarPipid && KstarKpid;

  //pt cuts
  TCut D0Kpt = "D0K_PT>400.0";
  TCut KstarKpt = "KstarK_PT>300.0";
  TCut KstarPipt = "KstarPi_PT>300.0";
  
  TCut ptcut = D0Kpt && KstarKpt && KstarPipt;

  //vertex cuts
  TCut BvertexChi2 = "Bd_ENDVERTEX_CHI2/Bd_ENDVERTEX_NDOF<4.0";
  TCut DvertexChi2 = "D0_ENDVERTEX_CHI2/D0_ENDVERTEX_NDOF<5.0";
  TCut BandKstarVertex = "(D0_ENDVERTEX_Z - Kstar_ENDVERTEX_Z)/sqrt(D0_ENDVERTEX_ZERR + Kstar_ENDVERTEX_ZERR)>-2.0";
  TCut vertex = BvertexChi2 && DvertexChi2 && BandKstarVertex;

  //MINIPCHI2
  TCut KstarMINIP = "Kstar_MINIPCHI2>25.0";
  TCut D0MINIP = "D0_MINIPCHI2>4.0";
  TCut BdMINIP = "Bd_MINIPCHI2<9.0";
  TCut MINIPCHI2 = KstarMINIP && D0MINIP && BdMINIP;
  
  TCut sel = mass_D0 && mass_Kstar && pidcut && ptcut && vertex && MINIPCHI2;
  
  fChain->Draw(">>elist",sel,"entrylist");
  cout << "Made the TEntryList" << endl;
  TEntryList *elist = (TEntryList*)gDirectory->Get("elist");

  fChain->SetEntryList(elist);
  cout << "applied event list to chain" << endl;

  TTree *small = fChain->CopyTree("");
  cout << "Events in selected tree (Before cos(helicity) and cos(Dir. angle w.r.t ownPV) cuts: " << small->GetEntries() << endl;

  coshel_and_Bdcosptng(small);
  cout << "added the vars to the small tree" << endl;
  
  //Bdcosptng and helicity cuts
  TCut cosptng = "Bdcosptng>0.99995";
  TCut coshelicity_cut = "abs(Kstar_coshel)>0.4";
  TCut final_cut = cosptng && coshelicity_cut;

  cout << "defined the cuts" << endl;

  //THIS IS THE LINE THAT CAUSES THE SEG FAULT  
  small->Draw(">>elist_final",final_cut);
  cout << "Made the FINAL TEventlistList" << endl;
  TEventList *elist1 = (TEventList*)gDirectory->Get("elist_final");

  TFile *ef = new TFile("selected_TEventList.root","recreate");         
  elist1->Write();
  ef->Close();

  small->SetEventList(elist1);
  cout << "applied FINAL event list to chain" << endl;

  TTree *final = small->CopyTree("");
  
  TFile *outputfile = new TFile("BdDKst_selected_events.root","recreate");  
  final->Write();
  cout << "Events in final selected tree: " << final->GetEntries() << endl;
  outputfile->Close();
  
}