#define Wpolarization_cxx
// The class definition in Wpolarization.h has been generated automatically
// by the Root utility TTree::MakeSelector.
//
// This class is derived from the Root class TSelector.
// The following members functions are called by the TTree::Process functions.
//    Begin:       called everytime a loop on the tree starts.
//                 a convenient place to create your histograms.
//    Notify():    This function is called at the first entry of a new Tree
//                 in a chain.
//    ProcessCut:  called at the beginning of each entry to return a flag
//                 true if the entry must be analyzed.
//    ProcessFill: called in the entry loop for all entries accepted 
//                 by Select.
//    Terminate:   called at the end of a loop on a TTree.
//                 a convenient place to draw/fit your histograms.
//
//   To use this file, try the following session on your Tree T
//
// Root > T.Process("Wpolarization.C")
// Root > T.Process("Wpolarization.C,"some options"")
// Root > T.Process("Wpolarization.C+")
//
#include "TH2.h"
#include "TStyle.h"
#include "TCanvas.h"
#include "TEventList.h"
#include "TLorentzVector.h"
#include <iostream.h>
#include "Wpolarization.h"

extern Float_t GetScaleww(Int_t energy,Char_t *channel);

TFile flist;

Int_t nevt = 0;
Float_t nexpected = 0;
Int_t nevqq = 0;
Int_t nmvqq = 0;
Int_t ntvqq = 0;
Int_t nwpcb = 0;
Int_t nwpnc = 0;
Int_t nkk2f = 0;
Int_t ndata = 0;
Int_t muon = 0;
Int_t electron = 0;
Bool_t useMuons,useElectrons,useList, fillList;

TEventList *allevents = 0;
TEventList *muonevents = 0;
TEventList *electronevents = 0;
TEventList *commonevents = 0;
TEventList *eventlist = 0;

TH1F *costwback;
TH1F *costwall;

TH1F *costhetastar_l;
TH1F *costhetastar_j;
TH1F *costhetastar_l_back;
TH1F *costhetastar_j_back;
TH1F *phistar_l;
TH1F *phistar_j;
TH1F *phistar_l_back;
TH1F *phistar_j_back;

TLorentzVector l_four_momentum;
TLorentzVector v_four_momentum;
TLorentzVector wl_four_momentum;
TLorentzVector wj_four_momentum;
TLorentzVector j1_four_momentum;
TLorentzVector j2_four_momentum;

void Wpolarization::Begin(TTree *tree)
{
// function called before starting the event loop
// initialize the tree branches

   Init(tree);

   TString option = GetOption();
   printf("Starting Wpolarization with process option: %s\n",option.Data());
   //
   // Some cleanup
   //
   gDirectory->Delete("costwback");
   gDirectory->Delete("costwall");
   gDirectory->Delete("costhetastar_l");
   gDirectory->Delete("costhetastar_j");
   gDirectory->Delete("costhetastar_l_back");
   gDirectory->Delete("costhetastar_j_back");
   gDirectory->Delete("phistar_l");
   gDirectory->Delete("phistar_j");
   gDirectory->Delete("phistar_l_back");
   gDirectory->Delete("phistar_j_back");

   // Creates the histograms;
   costwback = new TH1F("costwback","Cos #theta W", 11,-1,1);
   costwall = new TH1F("costwall","Cos Theta W",    11,-1,1);
   costhetastar_l = new TH1F("costhetastar_l","Cos #theta^{#star}_l",11,-1,1);
   costhetastar_j = new TH1F("costhetastar_j","Cos #theta^{#star}_j",11,-1,1);
   costhetastar_l_back = new TH1F("costhetastar_l_back","Cos #theta^{#star}_l",11,-1,1);
   costhetastar_j_back = new TH1F("costhetastar_j_back","Cos #theta^{#star}_j",11,-1,1);
   phistar_l = new TH1F("phistar_l","Phi^{#star}_l",11,-TMath::Pi(),TMath::Pi());
   phistar_j = new TH1F("phistar_j","Phi^{#star}_j",11,-TMath::Pi(),TMath::Pi());
   phistar_l_back = new TH1F("phistar_l_back","Phi^{#star}_l",11,-TMath::Pi(),TMath::Pi());
   phistar_j_back = new TH1F("phistar_j_back","Phi^{#star}_j",11,-TMath::Pi(),TMath::Pi());

   //process cases with event lists. First set it to false, then check
   // the options.
   fillList = kFALSE;
   useList  = kFALSE;
   useMuons = kFALSE;
   useElectrons = kFALSE;
   fChain->SetEventList(0);
   //
   // tries to open the file with the eventlists.

   if (option.Contains("useMuons")) {
      TFile flist("Weventlists.root");
      useList = kTRUE;
      eventlist = (TEventList*)flist.Get("muonevents");
      if (eventlist) eventlist->SetDirectory(0);
   }
   else if (option.Contains("useElectrons")) {
      TFile flist("Weventlists.root");
      useList = kTRUE;
      eventlist = (TEventList*)flist.Get("electronevents");
      if (eventlist) eventlist->SetDirectory(0);
   }
   else if (option.Contains("useList")) {
      TFile flist("Weventlists.root");
      useList = kTRUE;
      eventlist = (TEventList*)flist.Get("allevents");
      if (eventlist) eventlist->SetDirectory(0);
   }
   else {
      fillList = kTRUE;
   // Creates the Event lists 
      allevents      = new TEventList("allevents","Selection from Cut",50000);
      muonevents     = new TEventList("muonevents","mvqq",    25000);
      electronevents = new TEventList("electronevents","evqq",25000);
      commonevents   = new TEventList("commonevents","both mv ev",5000);
   }     
   if (useList) fChain->SetEventList(eventlist);

   // close the file with the event lists
   if (useList) flist.Close();
}

Bool_t Wpolarization::ProcessCut(Int_t entry)
{
// Selection function
// entry is the entry number in the current Tree
// Read only the necessary branches to select entries.
// return as soon as a bad entry is detected.
// to read complete event, call fChain->GetTree()->GetEntry(entry)
   muon = 0;
   electron = 0;

   if (useList) return kTRUE;

   fChain->GetTree()->GetEntry(entry);
   Float_t pmax = 0;

   for (Int_t i=0;i<Nne;i++) {
       if (Pne[i][3] >= 40 && Ipne[i]!=8 ) return kFALSE;
   }

   if (Impl>0) {
     if ((Pch[Impl-1][8] > 10)                      // muon ID
         && (TMath::Abs(Pch[Impl-1][3]-60) <= 40)   // lepton momentum
         && (Pch[Impl-1][17] > 10)                  // isolation angle
         && (TMath::Abs(Qm2ccm-80) <= 40)           // fitted W mass  
	 ) muon = 1;
     pmax = 0;
     for (Int_t i=0;i<Nch;i++) {
       if (i==Impl-1) continue;
       pmax = TMath::Max(pmax,Pch[i][3]);
     }
     if (pmax > 20) muon = 0;
   }

   if (Iepl>0) {
     if (  (Pch[Iepl-1][9] == 2)                     // electron ID
	   &&(TMath::Abs(Pch[Iepl-1][3]-60) <= 40)   // lepton momentum
	   &&(Pch[Iepl-1][17] > 20)                  // isolation angle
	   &&(Pmiss[3] > 20)                         // missing momentum
	   &&(TMath::Abs(Pmiss[2]/Pmiss[3])<0.8)     // theta missing mom.
	   && (TMath::Abs(Qe2ccm-80) <= 40)           // fitted W mass
	   ) electron = 1;
     pmax = 0;
     for (Int_t i=0;i<Nch;i++) {
       if (i==Iepl-1) continue;
       pmax = TMath::Max(pmax,Pch[i][3]);
     }
     if (pmax > 20) electron = 0;
   }

   if (muon == 0 && electron == 0) return kFALSE;

   if (fillList) {
      if (allevents) allevents->Enter(entry);
      if (muonevents && muon == 1) muonevents->Enter(entry);
      if (electronevents && electron == 1) electronevents->Enter(entry);
      if (commonevents && muon == 1 && electron == 1) commonevents->Enter(entry);
   }

   return kTRUE;
}

void Wpolarization::ProcessFill(Int_t entry)
{
// function called for selected entries only
// entry is the entry number in the current Tree
// read branches not processed in ProcessCut and fill histograms
// to read complete event, call fChain->GetTree()->GetEntry(entry)

   Char_t channel[5];
   Int_t Ilpl = 0; 
   Float_t Ql2cp[4][4];
   nevt++;

   if (useList) {
     fChain->GetEntry(entry);
     if (Impl >0) muon = 1;
     if (Iepl >0) electron = 1;
   }

   // Protection
   if (muon == 0 && electron == 0) {
     printf("\nERROR: Leptons index == 0\n");
     exit(-1);
   }

   //
   // Set up the channel name
   //
   // It should be:
   // evqq,mvqq,tvqq 
   // OR
   // wpcb -> background from wphactCC
   // wpnc -> background from wpahctNC
   // kk2f -> background from kk2f
   //
   if (Nrun < 0) {
     if ((Igen&1<<20) != 0) {
       sprintf(channel,"evqq");
     }
     else if ((Igen&1<<21) != 0) {
       sprintf(channel,"mvqq");
     }
     else if ((Igen&1<<22) != 0) {
       sprintf(channel,"tvqq");
     }
     else if ((Igen&1<<4) != 0) {
       sprintf(channel,"wpnc");
     }
     else if ((Igen&1<<1) != 0 || (Igen&1<<3) != 0) {
       sprintf(channel,"wpcb");
     }
     else if ((Igen&1<<5) != 0) {
       sprintf(channel,"kk2f");
     }
     else {
       printf("Unknown channel\n");
       exit(-1);
     }
   }
   else {
     sprintf(channel,"data");
   }

     
   // Copy the fitted matrices to the generic one

   if (muon == 1) {
     Ilpl = Impl;
     for(Int_t i=0;i<4;i++) {
       for(Int_t j =0;j<4;j++) {
	 Ql2cp[i][j] = Qm2cp[i][j];
       }
     }
   }
   if (electron == 1) {
     Ilpl = Iepl;
     for(Int_t i=0;i<4;i++) {
       for(Int_t j =0;j<4;j++) {
	 Ql2cp[i][j] = Qe2cp[i][j];
       }
     }
   }

   // Define the four vector components
   l_four_momentum.SetPx(Ql2cp[2][0]);
   l_four_momentum.SetPy(Ql2cp[2][1]);
   l_four_momentum.SetPz(Ql2cp[2][2]);
   l_four_momentum.SetE( Ql2cp[2][3]);
   v_four_momentum.SetPx(Ql2cp[3][0]);
   v_four_momentum.SetPy(Ql2cp[3][1]);
   v_four_momentum.SetPz(Ql2cp[3][2]);
   v_four_momentum.SetE( Ql2cp[3][3]);
   j1_four_momentum.SetPx(Ql2cp[0][0]);
   j1_four_momentum.SetPy(Ql2cp[0][1]);
   j1_four_momentum.SetPz(Ql2cp[0][2]);
   j1_four_momentum.SetE( Ql2cp[0][3]);
   j2_four_momentum.SetPx(Ql2cp[1][0]);
   j2_four_momentum.SetPy(Ql2cp[1][1]);
   j2_four_momentum.SetPz(Ql2cp[1][2]);
   j2_four_momentum.SetE( Ql2cp[1][3]);


   // calculates the W's four vectors
   wl_four_momentum  =l_four_momentum+v_four_momentum;
   wj_four_momentum  =j1_four_momentum+j2_four_momentum;

   // Apply the boost on the lepton and jet momentum. to get the decaying angle on the
   // W reference frame
   TVector3 boost =-((l_four_momentum+v_four_momentum)).BoostVector();
   TVector3 boostj = -((j1_four_momentum+j2_four_momentum)).BoostVector();

   //
   //
   //                   ATTENTION
   //
   // From now on, the charged lepton and the j1 four momentum is defined on the
   // W reference frame
   //
   //
   //
   l_four_momentum.Boost(boost);      // "boosted" lepton momentum
   j1_four_momentum.Boost(boostj);    // "boosted" jet 1 momentum
   if (Pch[Ilpl-1][5] == -1) {
     l_four_momentum.RotateY(-wl_four_momentum.Theta());
     j1_four_momentum.RotateY(-wj_four_momentum.Theta());
   }
   else {
     l_four_momentum.RotateY(-wl_four_momentum.Theta());
     j1_four_momentum.RotateY(-wj_four_momentum.Theta());
   }

   // Counts the number of events in each channel
   if (!strcmp(channel,"evqq")) nevqq++;
   if (!strcmp(channel,"mvqq")) nmvqq++;
   if (!strcmp(channel,"tvqq")) ntvqq++;
   if (!strcmp(channel,"wpcb")) nwpcb++;
   if (!strcmp(channel,"wpnc")) nwpnc++;
   if (!strcmp(channel,"kk2f")) nkk2f++;
   if (!strcmp(channel,"data")) ndata++;

   Float_t scale = 1.0;
   if (strcmp(channel,"data")) scale = GetScaleww(206,channel);

   // Plots the histograms for signal and background
   if (((Igen &1 << 21) == 0)&&((Igen &1 << 20) == 0)&&(Nrun<0)){   // neither a muon nor electron
     costwback->Fill(wl_four_momentum.CosTheta()*(-Pch[Ilpl-1][5]),scale);
     costwback->Fill(wj_four_momentum.CosTheta()*( Pch[Ilpl-1][5]),scale);
     costhetastar_l_back->Fill(l_four_momentum.CosTheta(),scale);
     costhetastar_j_back->Fill(j1_four_momentum.CosTheta(),scale);
     Float_t phistar = (1+Pch[Ilpl-1][5])*TMath::Pi()/2+l_four_momentum.Phi();
     if (phistar > TMath::Pi()) phistar-=2*TMath::Pi();
     phistar_l_back->Fill(phistar,scale);
     phistar_j_back->Fill(j1_four_momentum.Phi(),scale);
   }
   costwall->Fill(wl_four_momentum.CosTheta()*(-Pch[Ilpl-1][5]),scale);
   costwall->Fill(wj_four_momentum.CosTheta()*( Pch[Ilpl-1][5]),scale);
   costhetastar_l->Fill(l_four_momentum.CosTheta(),scale);
   costhetastar_j->Fill(j1_four_momentum.CosTheta(),scale);

   Float_t phistar = (1+Pch[Ilpl-1][5])*TMath::Pi()/2+l_four_momentum.Phi();
   if (phistar > TMath::Pi()) phistar-=2*TMath::Pi();
   phistar_l->Fill(phistar,scale);
   phistar_j->Fill(j1_four_momentum.Phi(),scale);

   nexpected+=scale;
}

void Wpolarization::Terminate()
{
// save the event list to a Root file if one was produced
   if (fillList) {
      TFile flist("Weventlists.root","recreate");
      allevents->Write();
      muonevents->Write();
      electronevents->Write();
      commonevents->Write();
      flist.Close();
   }
// function called at the end of the event loop

   cout << "Number of selected events: " << nevt  << "\t(" << nexpected << ")" << endl;
   cout << "Number of  evqq    events: " << nevqq << "\t(" << nevqq*GetScaleww(206,"evqq") << ")" << endl; 
   cout << "Number of  mvqq    events: " << nmvqq << "\t(" << nmvqq*GetScaleww(206,"mvqq") << ")" << endl;
   cout << "Number of  tvqq    events: " << ntvqq << "\t(" << ntvqq*GetScaleww(206,"tvqq") << ")" << endl;
   cout << "Number of  wpcb    events: " << nwpcb << "\t(" << nwpcb*GetScaleww(206,"wpcb") << ")" << endl;
   cout << "Number of  wpnc    events: " << nwpnc << "\t(" << nwpnc*GetScaleww(206,"wpnc") << ")" << endl;
   cout << "Number of  kk2f    events: " << nkk2f << "\t(" << nkk2f*GetScaleww(206,"kk2f") << ")" << endl;
   cout << "\nNumber of  data    events: " << ndata << endl;
}