#include "AnalyzeTree.h"
#include <TH2D.h>
#include <TH2.h>
#include <TLorentzVector.h>
#include <iostream>
#include <vector>
#include <math.h>
#include <string>
#include "myHist.C"
using namespace std;

// track quality cuts
static const int nTrackCut(6) ;
static const Int_t ndof_cut(5) ;
static const Int_t nchi2_cut(2) ;
static const Double_t dxy_cut(1.) ;
static const Double_t pt_cut(0.1) ;
static const Double_t eta_cut(2.5) ;
// split track cuts
static const Double_t cst_split_cut(0.9995) ;
static const Double_t dpt_split_cut(0.20) ;

static const Double_t pimass(0.1396), kamass(0.4937) ; 

vector <bool> GoodTrack ; // list of good tracks
vector <bool> SplitTrack ; // list of split tracks
vector <TLorentzVector> Tk4V ; // track four vectors (pi mass hyp.)

AnalyzeTree::AnalyzeTree() 
{}

AnalyzeTree::~AnalyzeTree() 
{}

// =======================================================================
void AnalyzeTree::analyze( int entry ) 
{
  if( numberOfTrack > 150 ) return ;   // sanity cut 
  SelectGoodTracks() ;  // as defined in GoodTrackNoDofCut
  RemoveSplitTracks() ; 
  ComputeBoseEinstein() ; 
  return;
}


// =======================================================================
// Compute B.E. correlations for the siganl and some control samples
// =======================================================================
void AnalyzeTree::ComputeBoseEinstein()
{ 
  for( int i = 1 ; i < numberOfTrack ; i++ )
    {
      if( !GoodTrack.at(i) ) continue ;
      TLorentzVector pv_i = Tk4V.at(i) ;
     
      for( int j = i+1 ; j< numberOfTrack ; j++ ) 
	{
	  if( !GoodTrack.at(j) ) continue ;
	  TLorentzVector pv_j = Tk4V.at(j) ;
	  double q = GetQ( &pv_i, &pv_j ) ;
	  
	  // invert II 4 vector (control sample) 
	  TLorentzVector j_vp( -pv_j.Px(), -pv_j.Py(), -pv_j.Pz() , pv_j.Energy() ) ;
	  double qinv = GetQ( &pv_i, &j_vp ) ;
	  
	  if( charge->at(i) != charge->at(j) ) 
	    { Hfill((char*)"qmix",q,100,0.,4.,1. ) ;}   // control sample +-
	  else if( charge->at(i)>0 ) 
	    { 
	      Hfill((char*)"qpp",q,100,0.,4.,1. ) ;     // signal ++
	      Hfill((char*)"qinvpp",qinv,100,0.,4.,1. ) ;  // control sample ++
	    }    
	  else if( charge->at(i)<0 ) 
	    { 
	      Hfill((char*)"qmm",q,100,0.,4.,1. ) ;    // signal -- 
	      Hfill((char*)"qinvmm",qinv,100,0.,4.,1. ) ;  // control sample ++
	    }
	}
    }
  return ;
}

// =======================================================================
// Compute correlated variable Q (= sqrt(qsq))
// =======================================================================
double AnalyzeTree::GetQ( TLorentzVector *p1, TLorentzVector *p2 )
{
  TLorentzVector Sum4V = (*p1)+(*p2) ;
  double q = Sum4V.Mag2() - 4*pimass*pimass  ;
  return (  q>0 ?  sqrt(q) : -sqrt(-q)  ) ; 
}

// =======================================================================
// Select good track, all cuts but ndof
// =======================================================================
inline void AnalyzeTree::SelectGoodTracks()
{
  Tk4V.erase( Tk4V.begin(), Tk4V.end() ) ;
  TLorentzVector p ;
  for ( int i = 0; i < numberOfTrack ; i++ ) 
    {
      p.SetXYZM( px->at(i), py->at(i), pz->at(i), pimass ) ;
      Tk4V.push_back( p ) ;

      GoodTrack.push_back(false) ;          
      if( !GoodTrackNoDofCut(i) ) continue ; // apply all cuts but that on NDof
      GoodTrack[i] = true ;

      if( charge->at(i) > 0 ) 
	{ 
	  Hfill( "EtaVsPtGoodP", pt->at(i), 100,0.,2.,eta->at(i), 60,-3.,3.,1. ) ; 
	  if( ndof->at(i) > ndof_cut ) Hfill( "EtaVsPtGoodP_ndof", pt->at(i), 100,0.,2.,eta->at(i), 60,-3.,3.,1. ) ; 
	}
      else
	{ 
	  Hfill( "EtaVsPtGoodM", pt->at(i), 100,0.,2.,eta->at(i), 60,-3.,3.,1. ) ;
	  if( ndof->at(i) > ndof_cut ) Hfill( "EtaVsPtGoodP_ndof", pt->at(i), 100,0.,2.,eta->at(i), 60,-3.,3.,1. ) ; 
	}
    } 
  return ;
}

// =======================================================================
// remove tracks having cos(theta)>cst_split_cut and 
// abs(delta p) < dpt_splt_cut  with another equal charge track 
// fill also some control histo
// =======================================================================
inline void  AnalyzeTree::RemoveSplitTracks() 
{
  for( int i = 0 ; i<numberOfTrack ; i++ ) 
    {
      SplitTrack.push_back( false ) ;
      if( !GoodTrack.at(i) ) continue ;
      TLorentzVector pv_i = Tk4V.at(i) ;

      for( int j = i+1 ; j<numberOfTrack ; j++ )
	{
	  if( !GoodTrack.at(j) ) continue ;

	  TLorentzVector pv_j = Tk4V.at(j) ;
	  double dpt = pv_i.Pt()-pv_j.Pt() ;
	  double cst = cos( pv_i.Angle( pv_j.Vect() ) ) ;	 
	  double q = GetQ( &pv_i, &pv_j ) ;

	  if( charge->at(i) != charge->at(j) ) 
	    {
	      Hfill("DptVsCst_pm",cst,50,0.9991,1.0001, dpt,50,-.5,.5,1.) ;
	      continue ;
	    }

	  // cos theta vs delta pt (/pt) for all charged tracks
	  string htit = "split" ;
	  if( charge->at(i) < 0 ) { htit.append("_m") ; }
	  Hfill((char*) htit.c_str(),cst,50,0.9991,1.0001, dpt,50,-.5,.5,1.) ;
	  htit.append("R") ;
	  Hfill((char*) htit.c_str(),cst,50,0.9991,1.0001, 2*dpt/(pv_i.Pt()+pv_j.Pt()),50,-1.,1.,1.) ;

	  if( cst >cst_split_cut && fabs(dpt)< dpt_split_cut ) 
	    { 
	      GoodTrack[j] = false ; 
	      SplitTrack[j] = true ;
	      Hfill("q_split",q,100,0.,4.,1. ) ;
	      Hfill("ndof_split",ndof->at(j),50,0.,50.,1.) ;
	      Hfill("eta_vs_pt_split",pt->at(j),50,0.,1.,eta->at(j),60,-3.,3.,1.) ;
	      Hfill("eta_vs_phi_split",phi->at(j),65,-3.25,3.25,eta->at(j),60,-3.,3.,1.) ;
	      Hfill("eta_vs_dz_split",dz->at(j),50,-5.,5.,eta->at(j),60,-3.,3.,1.) ;
	      Hfill("eta_vs_ddz_split",dz->at(j)-dz->at(i),100,-5.,5.,eta->at(j),60,-3.,3.,1.) ;
	    }	  
	}
    }
  return ;
}
// =======================================================================
inline void AnalyzeTree::SaveToFile( char* OutputFileName ) { SaveHist( OutputFileName ) ;}
// =======================================================================
inline void AnalyzeTree::end_analysis() {}
// =======================================================================
inline void AnalyzeTree::begin_analysis() { return ; }
// =======================================================================
inline int AnalyzeTree::GoodTrackNoDofCut( int i ) 
{
  if( RejectTrack(i) == 0 || RejectTrack(i) == (int) pow(2,4) ) return 1 ;
  return 0 ;
}
// =======================================================================
inline int AnalyzeTree::RejectTrack( int i ) 
{
  int ReturnValue(0) ;

  if( !highPurity_track->at(i) )         ReturnValue += (int) pow(2,0) ;
  if( pt->at(i) < pt_cut )               ReturnValue += (int) pow(2,1) ;
  if( fabs( dxy_PV->at(i)) > dxy_cut )   ReturnValue += (int) pow(2,2) ;
  if( normalizedChi2->at(i) > nchi2_cut) ReturnValue += (int) pow(2,3) ;
  if( ndof->at(i) < ndof_cut )           ReturnValue += (int) pow(2,4) ;
  if( fabs( eta->at(i) ) > eta_cut )     ReturnValue += (int) pow(2,5) ;

  FillCutControlHisto( i, ReturnValue ) ;
 
  return ReturnValue ;
}

// =======================================================================
inline void AnalyzeTree::FillCutControlHisto( int itrack, int rcut )
{
  char htitle[100] ;
  for( int jcut=0 ; jcut<nTrackCut ; jcut++ ) 
    { 
      if( (rcut&(int) pow(2,jcut) ) ) continue ;
      int h = sprintf( htitle, "etacut%d", jcut ) ;
      if( charge->at(itrack) > 0 ) 
	{ strcat( htitle,"_plus" ) ; }
      else
	{ strcat( htitle,"_minus" ) ; }
      Hfill( htitle, eta->at(itrack), 60,-3.,3.,1. ) ;      
    }
}