//----------------------------------------------------
//
// TLikelihoodLBE - identifies leptons using likelihood method
//    - calculate likelihood of the object being an electron (central)
//      according provided templates for real electrons and fakes
//   
// author: lysak@fnal.gov
// modified for topntuple usage -- w yao (9/30/2010)
//----------------------------------------------------
#include<iostream>

#include "TFile.h"
#include "TTree.h"
#include "TH1D.h"

#include "TLikelihoodLBE.hh"

using namespace std;

//ClassImp(TLikelihoodLBE)


TLikelihoodLBE::TLikelihoodLBE():
  _inFileSignal (0),
  _inFileBckg   (0),
  _likelihood  (-1),
  _debug        (0),
  _inputIsFile  (false),
  _inputIsHist  (false)
{
  for (int i=0; i<kNumVariables; i++)  {
    _hVarSignal[i]    = 0;
    _hVarBckg[i]      = 0;
    _usingVariable[i] = true;
  }
}


TLikelihoodLBE::TLikelihoodLBE(char* inFileSignal,
			       char* inFileBckg):
  _likelihood  (-1),
  _debug        (0),
  _inputIsFile  (true),
  _inputIsHist  (false)
{
  _inFileSignal = new TFile(inFileSignal);
  _inFileBckg   = new TFile(inFileBckg);

  for (int i=0; i< kNumVariables; i++)  {
    _hVarSignal[i]    = 0;
    _hVarBckg[i]      = 0;
    _usingVariable[i] = true;
  }

  //BE CAREFULL!!! if you 'Get' object from file, it will be added into memory
  //and you get the pointer to memory object, so you MUST delete this afterwards!!!
  _hVarSignal[kEt         ] = (TH1D*) _inFileSignal->Get("Et;1");  
  _hVarSignal[kEta        ] = (TH1D*) _inFileSignal->Get("Eta;1"); 
  _hVarSignal[kHadEm      ] = (TH1D*) _inFileSignal->Get("HadEm;1");
  _hVarSignal[kLshr       ] = (TH1D*) _inFileSignal->Get("Lshr;1");
  _hVarSignal[kTrackPt    ] = (TH1D*) _inFileSignal->Get("TrackPt;1");
  _hVarSignal[kTrackZ0    ] = (TH1D*) _inFileSignal->Get("TrackZ0;1");
  _hVarSignal[kZ0         ] = (TH1D*) _inFileSignal->Get("Z0;1");
  _hVarSignal[kNCotHitsTot] = (TH1D*) _inFileSignal->Get("NCotHitsTot;1");
  _hVarSignal[kNCotHitsAx ] = (TH1D*) _inFileSignal->Get("NCotHitsAx;1");
  _hVarSignal[kNCotHitsSt ] = (TH1D*) _inFileSignal->Get("NCotHitsSt;1");
  _hVarSignal[kNSvxHits   ] = (TH1D*) _inFileSignal->Get("NSvxHits;1");
  _hVarSignal[kNCotAxSeg  ] = (TH1D*) _inFileSignal->Get("NCotAxSeg;1");
  _hVarSignal[kNCotStSeg  ] = (TH1D*) _inFileSignal->Get("NCotStSeg;1");
  _hVarSignal[kChi2Cot    ] = (TH1D*) _inFileSignal->Get("Chi2Cot;1");
  _hVarSignal[kEOverP     ] = (TH1D*) _inFileSignal->Get("EOverP;1");
  _hVarSignal[kQDelX      ] = (TH1D*) _inFileSignal->Get("QDelX;1");
  _hVarSignal[kDelZ       ] = (TH1D*) _inFileSignal->Get("DelZ;1");
  _hVarSignal[kXCes       ] = (TH1D*) _inFileSignal->Get("XCes;1");
  _hVarSignal[kZCes       ] = (TH1D*) _inFileSignal->Get("ZCes;1");
  _hVarSignal[kCalIso     ] = (TH1D*) _inFileSignal->Get("CalIso;1");
  _hVarSignal[kTrkIso     ] = (TH1D*) _inFileSignal->Get("TrkIso;1");

  _hVarBckg[kEt         ] = (TH1D*) _inFileBckg->Get("Et;1");  
  _hVarBckg[kEta        ] = (TH1D*) _inFileBckg->Get("Eta;1"); 
  _hVarBckg[kHadEm      ] = (TH1D*) _inFileBckg->Get("HadEm;1");
  _hVarBckg[kLshr       ] = (TH1D*) _inFileBckg->Get("Lshr;1");
  _hVarBckg[kTrackPt    ] = (TH1D*) _inFileBckg->Get("TrackPt;1");
  _hVarBckg[kTrackZ0    ] = (TH1D*) _inFileBckg->Get("TrackZ0;1");
  _hVarBckg[kZ0         ] = (TH1D*) _inFileBckg->Get("Z0;1");
  _hVarBckg[kNCotHitsTot] = (TH1D*) _inFileBckg->Get("NCotHitsTot;1");
  _hVarBckg[kNCotHitsAx ] = (TH1D*) _inFileBckg->Get("NCotHitsAx;1");
  _hVarBckg[kNCotHitsSt ] = (TH1D*) _inFileBckg->Get("NCotHitsSt;1");
  _hVarBckg[kNSvxHits   ] = (TH1D*) _inFileBckg->Get("NSvxHits;1");
  _hVarBckg[kNCotAxSeg  ] = (TH1D*) _inFileBckg->Get("NCotAxSeg;1");
  _hVarBckg[kNCotStSeg  ] = (TH1D*) _inFileBckg->Get("NCotStSeg;1");
  _hVarBckg[kChi2Cot    ] = (TH1D*) _inFileBckg->Get("Chi2Cot;1");
  _hVarBckg[kEOverP     ] = (TH1D*) _inFileBckg->Get("EOverP;1");
  _hVarBckg[kQDelX      ] = (TH1D*) _inFileBckg->Get("QDelX;1");
  _hVarBckg[kDelZ       ] = (TH1D*) _inFileBckg->Get("DelZ;1");
  _hVarBckg[kXCes       ] = (TH1D*) _inFileBckg->Get("XCes;1");
  _hVarBckg[kZCes       ] = (TH1D*) _inFileBckg->Get("ZCes;1");
  _hVarBckg[kCalIso     ] = (TH1D*) _inFileBckg->Get("CalIso;1");
  _hVarBckg[kTrkIso     ] = (TH1D*) _inFileBckg->Get("TrkIso;1");
}


TLikelihoodLBE::TLikelihoodLBE(TH1D** templates_signal, 
			       TH1D** templates_bckg):
  _inFileSignal (0),
  _inFileBckg   (0),
  _likelihood  (-1),
  _debug        (0),
  _inputIsFile  (false),
  _inputIsHist  (true)
{
  for (int i=0; i< kNumVariables; i++)   {
    _hVarSignal[i]    = 0;
    _hVarBckg[i]      = 0;
    _usingVariable[i] = true;
  }

  //basic check that the histograms exist
  if ( ! (templates_signal && templates_bckg) )
    return ;

  //BE CAREFULL!!! here we don't clone histograms (to be quicker)
  //so, you should not delete them at the end since you don't own them 
  //(they will be used later)!
  _hVarSignal[kEt         ] = (TH1D*) templates_signal[kEt];  
  _hVarSignal[kEta        ] = (TH1D*) templates_signal[kEta]; 
  _hVarSignal[kHadEm      ] = (TH1D*) templates_signal[kHadEm];
  _hVarSignal[kLshr       ] = (TH1D*) templates_signal[kLshr];
  _hVarSignal[kTrackPt    ] = (TH1D*) templates_signal[kTrackPt];
  _hVarSignal[kTrackZ0    ] = (TH1D*) templates_signal[kTrackZ0];
  _hVarSignal[kZ0         ] = (TH1D*) templates_signal[kZ0];
  _hVarSignal[kNCotHitsTot] = (TH1D*) templates_signal[kNCotHitsTot];
  _hVarSignal[kNCotHitsAx ] = (TH1D*) templates_signal[kNCotHitsAx];
  _hVarSignal[kNCotHitsSt ] = (TH1D*) templates_signal[kNCotHitsSt];
  _hVarSignal[kNSvxHits   ] = (TH1D*) templates_signal[kNSvxHits];
  _hVarSignal[kNCotAxSeg  ] = (TH1D*) templates_signal[kNCotAxSeg];
  _hVarSignal[kNCotStSeg  ] = (TH1D*) templates_signal[kNCotStSeg];
  _hVarSignal[kChi2Cot    ] = (TH1D*) templates_signal[kChi2Cot];
  _hVarSignal[kEOverP     ] = (TH1D*) templates_signal[kEOverP];
  _hVarSignal[kQDelX      ] = (TH1D*) templates_signal[kQDelX];
  _hVarSignal[kDelZ       ] = (TH1D*) templates_signal[kDelZ];
  _hVarSignal[kXCes       ] = (TH1D*) templates_signal[kXCes];
  _hVarSignal[kZCes       ] = (TH1D*) templates_signal[kZCes];
  _hVarSignal[kCalIso     ] = (TH1D*) templates_signal[kCalIso];
  _hVarSignal[kTrkIso     ] = (TH1D*) templates_signal[kTrkIso];

  _hVarBckg[kEt         ] = (TH1D*) templates_bckg[kEt];  
  _hVarBckg[kEta        ] = (TH1D*) templates_bckg[kEta]; 
  _hVarBckg[kHadEm      ] = (TH1D*) templates_bckg[kHadEm];
  _hVarBckg[kLshr       ] = (TH1D*) templates_bckg[kLshr];
  _hVarBckg[kTrackPt    ] = (TH1D*) templates_bckg[kTrackPt];
  _hVarBckg[kTrackZ0    ] = (TH1D*) templates_bckg[kTrackZ0];
  _hVarBckg[kZ0         ] = (TH1D*) templates_bckg[kZ0];
  _hVarBckg[kNCotHitsTot] = (TH1D*) templates_bckg[kNCotHitsTot];
  _hVarBckg[kNCotHitsAx ] = (TH1D*) templates_bckg[kNCotHitsAx];
  _hVarBckg[kNCotHitsSt ] = (TH1D*) templates_bckg[kNCotHitsSt];
  _hVarBckg[kNSvxHits   ] = (TH1D*) templates_bckg[kNSvxHits];
  _hVarBckg[kNCotAxSeg  ] = (TH1D*) templates_bckg[kNCotAxSeg];
  _hVarBckg[kNCotStSeg  ] = (TH1D*) templates_bckg[kNCotStSeg];
  _hVarBckg[kChi2Cot    ] = (TH1D*) templates_bckg[kChi2Cot];
  _hVarBckg[kEOverP     ] = (TH1D*) templates_bckg[kEOverP];
  _hVarBckg[kQDelX      ] = (TH1D*) templates_bckg[kQDelX];
  _hVarBckg[kDelZ       ] = (TH1D*) templates_bckg[kDelZ];
  _hVarBckg[kXCes       ] = (TH1D*) templates_bckg[kXCes];
  _hVarBckg[kZCes       ] = (TH1D*) templates_bckg[kZCes];
  _hVarBckg[kCalIso     ] = (TH1D*) templates_bckg[kCalIso];
  _hVarBckg[kTrkIso     ] = (TH1D*) templates_bckg[kTrkIso];
}


TLikelihoodLBE::~TLikelihoodLBE() 
{
  //delete depending on whether the templates were read 
  //from file or just from hists 
  if (_inputIsFile) {
    for (int i=0; i< kNumVariables; i++)
      {
	if ( _hVarSignal[i]) delete _hVarSignal[i];
	if ( _hVarBckg[i]  ) delete _hVarBckg[i];
      }
  }

  if ( (_inFileSignal) && (_inFileBckg) ) {
    _inFileBckg->Close();
    _inFileSignal->Close();
    
    delete _inFileBckg;
    delete _inFileSignal;
  }

}


int TLikelihoodLBE::SetAllVariablesUsage(bool useVar)
{ 
  for (int i=0; i<kNumVariables; i++) _usingVariable[i] = useVar ; 
  
  return 0;
}


int TLikelihoodLBE::CalculateLikelihood(const electron * ele, const offltrack* eletrk)
{
  int rc = ReadVariables(ele, eletrk);
  if (rc < 0) return -1;
  double likelihood  =1;
  double prob_signal =1;
  double prob_bckg   =1;

  for (int i=0; i<kNumVariables; i++)
    {
      if ( ! _usingVariable[i] ) continue;

      double p_s = Psignal((LepIDvariable)i) ;
      double p_b = Pbckg  ((LepIDvariable)i) ;
      if ( _debug > 3 ) 
	cout << "ID,value,Psig,Pbckg= " << i   << " " << variable[i] << " " 
	                                << p_s << " " << p_b         << endl;
      likelihood  *= p_s / (p_s + p_b) ; 
      prob_signal *= p_s ; 
      prob_bckg   *= p_b ;
    }

  //_likelihood = likelihood;
  _likelihood = (prob_signal)/(prob_signal + prob_bckg);
  
  if (_debug > 1) cout << "likelihood = " << _likelihood << endl;
  if (_debug > 2) cout << "signal,bckg. probab. = " << prob_signal << " " << prob_bckg << endl;
 
  return 0;
} 


int TLikelihoodLBE::ReadVariables(const electron * el, const offltrack * eletrk)
{
  //  
  if ( el->Region !=0 ) {
    cout << "TLikelihoodLBE::ReadVariables : electron is NOT an central electron!"<<endl;
    return -1;
  }
  if(!eletrk){
    cout<<" TLikelihoodLBE::ReadVariables : electron do not have offltrack linked "<<endl;
    return -1;
  }
  variable[kEt         ] = el->Et;
  variable[kEta        ] = el->Eta;
  variable[kHadEm      ] = el->Hadem;
  variable[kLshr       ] = el->LshrTrk; // ?
  variable[kTrackPt    ] = el->TrkPt;
  variable[kTrackZ0    ] = el->TrkZ0;
  variable[kZ0         ] = el->TrkZ0;
  variable[kNCotHitsTot] = el->TrkAxHits+el->TrkStHits;
  variable[kNCotHitsAx ] = el->TrkAxHits; 
  variable[kNCotHitsSt ] = el->TrkStHits;
  variable[kNSvxHits   ] = el->TrkSiHits;
  //
  int nslax(0);
  int nslst(0);
  for(int sl=0; sl<8; ++sl){
    if(eletrk->NumCTHits[sl]>=5){ // requiring 5 or more hits 
      if(sl%2==0) nslst++; 
      else nslax++;
    }
  }
  variable[kNCotAxSeg  ] = nslax; // el->TrkAxSeg; // for 7 or more hits 
  variable[kNCotStSeg  ] = nslst; // el->TrkStSeg;
  variable[kChi2Cot    ] = eletrk->chi2CT; 
  variable[kEOverP     ] = el->EP;
  variable[kQDelX      ] = el->DeltaX*el->Charge;
  variable[kDelZ       ] = el->DeltaZ;
  variable[kXCes       ] = el->TrkCESx;
  variable[kZCes       ] = el->TrkCESz;
  variable[kCalIso     ] = el->Isol;
  variable[kTrkIso     ] = el->TrkIsol;
  
  return 0;
}


double TLikelihoodLBE::Psignal(LepIDvariable var)
{
  return HistValue(_hVarSignal[var],variable[var]);

}


double TLikelihoodLBE::Pbckg(LepIDvariable var)
{
  return HistValue(_hVarBckg[var],variable[var]);
}


double TLikelihoodLBE::Psignal(LepIDvariable var,double value)
{
  return HistValue(_hVarSignal[var],value);
}


double TLikelihoodLBE::Pbckg(LepIDvariable var,double value)
{
  return HistValue(_hVarBckg[var],value);
}


double TLikelihoodLBE::HistValue(TH1D* hist, double value)
{
  //if the value is outside the range take underflow/overflow
  if ( value < hist->GetBinLowEdge(1))
    return hist->GetBinContent(0);
    //return hist->GetBinContent(0)*hist->GetBinWidth(0);

  int nbins = hist->GetNbinsX();
  if (value >= (hist->GetBinLowEdge(nbins) + hist->GetBinWidth(nbins)) )
    return hist->GetBinContent(nbins+1);
    //return hist->GetBinContent(nbins+1)*hist->GetBinWidth(nbins+1);


  //find the bin which the value belongs to 
  //assuming all bins have the same width && 
  //if value is at bin boundary, use the higher bin,i.e <a,b) interval used for a bin  
  double bin_width = hist->GetBinWidth(1);
  int    bin_id    = (int) ( (value - hist->GetBinLowEdge(1)) / bin_width ) + 1;
  
  if (_debug >= 10) cout << "bin = " << bin_id << endl;

  return hist->GetBinContent(bin_id);
}