#include <iostream>
#include <iomanip>
#include <cmath>
#include <Stntuple/obj/TStnElectronBlock.hh>
#include <Stntuple/obj/TStnMuonBlock.hh>
#include <Stntuple/obj/TStnPhotonBlock.hh>
#include "Stntuple/alg/TStntuple.hh"
#include "Stntuple/obj/TStnTrack.hh"

#include "Stntuple/alg/TStnLikelihood.hh"

using std::cout;
using std::endl;

ClassImp(TStnLikelihood)

//_____________________________________________________________________________
TStnLikelihood::TStnLikelihood()
{
}

//_____________________________________________________________________________
TStnLikelihood::~TStnLikelihood() {
}

//_____________________________________________________________________________
Int_t TStnLikelihood::PreSelCenEle(TStnElectron* ele) {
  int flag = 0;
  if(ele->Et()<20.0) flag |= (1<<0);
  if(ele->Et()<50.0 && ele->EOverP()>2.0) flag |= (1<<1);
  if(ele->FidEleSmx()!=1) flag |= (1<<2);
  if(fabs(ele->Z0())>60.0) flag |= (1<<3);
  if(ele->DelXQ()<-3.0 || ele->DelXQ()>1.5) flag |= (1<<4);
  if(fabs(ele->DelZ())>3.0) flag |= (1<<5);
  TStnTrack* trk = ele->Track();
  if(!trk) flag |= (1<<6);
  if(trk) {
    if(trk->NCotAxSeg(5)<3 || trk->NCotStSeg(5)<3) flag |= (1<<7);
  }
  return flag;
}

//_____________________________________________________________________________
Int_t TStnLikelihood::PreSelPlgEle(TStnElectron* phoenixEle) {
  int flag = 0;
  if(phoenixEle->Et()<20.0) flag |= (1<<0);
  if(fabs(phoenixEle->Z0())>60.0) flag |= (1<<1);
  if(phoenixEle->HadEm()>0.125) flag |= (1<<2);
  TStnTrack* trk = phoenixEle->Track();
  if(!trk) flag |= (1<<3);
  if(trk) {
    if(trk->NSvxHits()<=3) flag |= (1<<4);
  }
  if(fabs(phoenixEle->DetEta())<1.2 || 
     fabs(phoenixEle->DetEta())>2.5) flag |= (1<<5);
  return flag;
}

//_____________________________________________________________________________
Int_t TStnLikelihood::PreSelMuon(TStnMuon* muo, Float_t vertexZ) {
  int flag = 0;
  if(muo->HadEnergy()<0.0) flag |= (1<<0);
  if(!muo->HasCmuStub()) flag |= (1<<1);
  if(!muo->HasCmpStub()) flag |= (1<<2);
  float d0 = TStntuple::CorrectedD0(muo);
  if(fabs(d0)>2.0) flag |= (1<<3);
  //if(vertexZ>-998.0 && fabs(vertexZ-muo->Z0())>5.0) flag |= (1<<4);
  TStnTrack* trk = muo->Track();
  if(!trk) flag |= (1<<5);
  if(trk) {
    if(trk->NCotAxSeg(5)<3 || trk->NCotStSeg(5)<3) flag |= (1<<6);
  }
  return flag;
}

//_____________________________________________________________________________
Float_t TStnLikelihood::CenEle(TStnElectron* ele) {
  if(ele->TrackPt()<0.0 || ele->TrackPt()>1e30) return -1;
  double tIso = ele->TrackIso()*ele->TrackPt();
  if(tIso<0.0 || tIso>2000.0) return -1;
  double isoGeV = ele->IsoCorr();
  if(isoGeV<-100.0 || isoGeV>2000.0) return -1;
  double lshr = ele->Lshr();
  if(lshr<-100.0 || lshr>100.0) return -1;
  double hadem = ele->HadEm();
  if(hadem<0.0 || hadem>100.0) return -1;
  double cesStr = ele->Chi2Strip();
  if(cesStr<0.0 || cesStr>2000.0) return -1;
  double rlike = like.CentralEmIdLikelihood(tIso, isoGeV, lshr, hadem, cesStr);
  return Norm(rlike);
}

//_____________________________________________________________________________
Float_t TStnLikelihood::PlgEle(TStnElectron* ele) {
  double fraciso4 = ele->Iso1Corr();
  if(fraciso4<-100.0 || fraciso4>100.0) return -1;
  double pes = (ele->Pes5x9(0)<ele->Pes5x9(1) ? 
		ele->Pes5x9(0) : ele->Pes5x9(1) );
  if(pes<0.0 || pes>1.0) return -1;
  double hadem = ele->HadEm();
  if(hadem<0.0 || hadem>100.0) return -1;
  double rlike = like.PlugEmIdLikelihood( fraciso4, pes, hadem);
  return Norm(rlike);
}

/*
//_____________________________________________________________________________
Float_t TStnLikelihood::CenPho(TStnPhoton* pho) {
  double tIso = pho->Tiso();
  double isoGeV = pho->EIso4(1);
  double lshr = pho->LShr();
  double hadem = pho->HadEm();
  double cesStr = pho->Chi2Strip();
  double like = like.CentralEmIdLikelihood(tIso, isoGeV, lshr, hadem, cesStr);
  return Norm(like);
}

//_____________________________________________________________________________
Float_t TStnLikelihood::PlgPho(TStnPhoton* pho) {
  double fraciso4 = pho->EIso4(1);
  double pes = pho->Prof5by9U();
  double hadem = pho->HadEm();
  double rlike = like.PlugEmIdLikelihood( fraciso4, pes, hadem);
  return Norm(rlike);
}
*/

//_____________________________________________________________________________
Float_t TStnLikelihood::Muon(TStnMuon* muo) {

  double iso4 = muo->ConeR4Et();
  if(iso4<0.0 || iso4>2000.0) return -1;
  double tIso = muo->TIso()*muo->TrackPt();
  if(tIso<0.0 || tIso>2000.0) return -1;
  double hadE = muo->HadEnergy();
  if(hadE<0.0 || hadE>2000.0) return -1;
  double emE = muo->EmEnergy();
  if(emE<0.0 || emE>2000.0) return -1;
  double rlike = like.MuonLikelihood(iso4, tIso, hadE, emE);
  return Norm(rlike);

}

//_____________________________________________________________________________
Float_t TStnLikelihood::Norm(Double_t like) {
  if(like<-50 || like>50) return -1.0;
  like = exp(like);
  return like/(1.0+like);
}