#include "StHbtMaker/Cut/gregsProtonTrackCut.h"
#include <sstream>

#ifdef __ROOT__
ClassImp(gregsProtonTrackCut)
#endif

//_________________
gregsProtonTrackCut::gregsProtonTrackCut() { //Constructor

  mNTracksPassed = 0;
  mNTracksFailed = 0;

  mSelType = 3;

  mType = 1;
  mCharge = 0;
  mNHits[0] = 14; mNHits[1] = 46;
  mFlag[0] = 0; mFlag[1] = 1000;
  mAntiSplit = 0.;         //nHitsFit/nHitsPoss
  mNHitsFit[0] = 0; mNHitsFit[1] = 46;
  mP[0] = 0.15; mP[1] = 1.5;
  mPt[0] = -1e9; mPt[1] = +1e9;
  mPx[0] = -1e9; mPx[1] = +1e9;
  mPy[0] = -1e9; mPy[1] = +1e9;
  mPz[0] = -1e9; mPz[1] = +1e9;
  mRapidity[0] = -1e9; mRapidity[1] = +1e9;
  mEta[0] = -1.; mEta[1] = 1.;
  mDCA[0] = -0.01; mDCA[0] = 3.;
  mDCAGlobal[0] = -0.01; mDCAGlobal[1] = 3.;

  mTpcMomentum[0] = 0.15; mTpcMomentum[1] = 0.55;
  mTpcNSigmaElectron[0] = -2.; mTpcNSigmaElectron[1] = 2.;
  mTpcNSigmaPion[0] = -2.; mTpcNSigmaPion[1] = 2.;
  mTpcNSigmaKaon[0] = -2.; mTpcNSigmaKaon[1] = 2.;
  mTpcNSigmaProton[0] = -2.; mTpcNSigmaProton[1] = 2.;

  mTofMomentum[0] = 0.15; mTofMomentum[1] = 1.45;
  mTofMassSqr[0] = 0.2; mTofMassSqr[1] = 0.32;
  mTofInvBetaDiff[0] = -1e9; mTofInvBetaDiff[1] = +1e9;

  mTnTMomentum[0] = 0.15; mTnTMomentum[1] = 1.45;
  mTnTNSigmaProton[0] = -3.; mTnTNSigmaProton[1] = 3.;
}

//_________________
gregsProtonTrackCut::gregsProtonTrackCut(const gregsProtonTrackCut &c) : StHbtTrackCut(c) {

  mNTracksPassed = 0;
  mNTracksFailed = 0;

  mSelType = c.mSelType;

  mCharge = c.mCharge;
  mType = c.mType;
  mAntiSplit = c.mAntiSplit;
  mFlag[0] = c.mFlag[0];          mFlag[1] = c.mFlag[1];
  mNHits[0] = c.mNHits[0];        mNHits[1] = c.mNHits[1];
  mNHitsFit[0] = c.mNHitsFit[0];  mNHitsFit[1] = c.mNHitsFit[1];
  mP[0] = c.mP[0];   mP[1] = c.mP[1];
  mPt[0] = c.mPt[0]; mPt[1] = c.mPt[1];
  mPx[0] = c.mPx[0]; mPx[1] = c.mPx[1];
  mPy[0] = c.mPy[0]; mPy[1] = c.mPy[1];
  mPz[0] = c.mPz[0]; mPz[1] = c.mPz[1];
  mRapidity[0] = c.mRapidity[0]; mRapidity[1] = c.mRapidity[1];
  mEta[0] = c.mEta[0]; mEta[1] = c.mEta[1];
  mDCA[0] = c.mDCA[0]; mDCA[1] = c.mDCA[1];
  mDCAGlobal[0] = c.mDCAGlobal[0]; mDCAGlobal[1] = c.mDCAGlobal[1];

  mTpcMomentum[0] = c.mTpcMomentum[0];             mTpcMomentum[1] = c.mTpcMomentum[1];
  mTpcNSigmaElectron[0] = c.mTpcNSigmaElectron[0]; mTpcNSigmaElectron[1] = c.mTpcNSigmaElectron[1];
  mTpcNSigmaPion[0] = c.mTpcNSigmaPion[0];         mTpcNSigmaPion[1] = c.mTpcNSigmaPion[1];
  mTpcNSigmaKaon[0] = c.mTpcNSigmaKaon[0];         mTpcNSigmaKaon[1] = c.mTpcNSigmaKaon[1];
  mTpcNSigmaProton[0] = -c.mTpcNSigmaProton[0];    mTpcNSigmaProton[1] = c.mTpcNSigmaProton[1];

  mTofMomentum[0] = c.mTofMomentum[0];       mTofMomentum[1] = c.mTofMomentum[1];
  mTofMassSqr[0] = c.mTofMassSqr[0];         mTofMassSqr[1] = c.mTofMassSqr[1];
  mTofInvBetaDiff[0] = c.mTofInvBetaDiff[0]; mTofInvBetaDiff[1] = c.mTofInvBetaDiff[1];

  mTnTMomentum[0] = c.mTnTMomentum[0]; mTnTMomentum[1] = c.mTnTMomentum[1];
  mTnTNSigmaProton[0] = c.mTnTNSigmaProton[0]; mTnTNSigmaProton[1] = c.mTnTNSigmaProton[1];
}

//_________________
gregsProtonTrackCut::~gregsProtonTrackCut() {
  /* empty */
}

//_________________
bool gregsProtonTrackCut::Pass(const StHbtTrack *track) {

  float mEnergy = ::sqrt(track->P().mag2() + mMass*mMass);
  float mY = 0.5*::log( (mEnergy+track->P().z()) / (mEnergy-track->P().z()) );

  bool mGoodCharge = false;
  bool mGoodTrack = false;
  bool mGoodPID = false;
  bool mIsTofTrack = false;
  float mom = track->P().mag();
  
  bool mPassTrack = false;

  //Check the correctness of the detector selection
  if(mSelType<0 || mSelType>3) {
    std::cout << "gregsProtonTrackCut::Pass [WARNING] Wrong detector selection type"
	      << mSelType << " . Resetting to 3" << std::endl;
    mSelType = 3;
  }

  //Is the right charge
  if(track->Charge() == mCharge) {
    mGoodCharge = true;
  }

#ifdef STHBTDEBUG
  std::cout << "gregsProtonTrackCut [DEBUG]" << std::endl
	    << "type: " << track->TrackType() 
	    << " nHF/nHPoss: " << track->NHitsFit2PossRatio()
	    << " p: " << track->P().mag() << " eta: " << track->Eta()
	    << " y: " << mY << " beta: " << track->TofBeta() 
	    << " flag: " << track->Flag() <<
	    << " dcaXYGlobal: " << track->DCAxyGlobal() << std::endl;
#endif

  //Is the good track
  if( (track->TrackType() == mType) &&
      (track->NHitsFit2PossRatio() >= mAntiSplit) &&
      (track->DCAxy() >= mDCA[0]) && 
      (track->DCAxy() <= mDCA[1]) &&
      (track->DCAxyGlobal() >= mDCAGlobal[0]) && 
      (track->DCAxyGlobal() <= mDCAGlobal[1]) &&
      (track->NHits() >= mNHits[0]) && 
      (track->NHits() <= mNHits[1]) &&
      (track->NHitsFit() >= mNHitsFit[0]) && 
      (track->NHitsFit() <= mNHitsFit[1]) &&
      (mom >= mP[0]) && 
      (mom <= mP[1]) &&
      (track->Pt() >= mPt[0]) && 
      (track->Pt() <= mPt[1]) &&
      (track->P().x() >= mPx[0]) && 
      (track->P().x() <= mPx[1]) &&
      (track->P().y() >= mPy[0]) && 
      (track->P().y() <= mPy[1]) &&
      (track->P().z() >= mPz[0]) && 
      (track->P().z() <= mPz[1]) &&
      (track->Eta() >= mEta[0]) && 
      (track->Eta() <= mEta[1]) &&
      (mY >= mRapidity[0]) && 
      (mY <= mRapidity[1]) &&
      (track->Flag() >= mFlag[0]) && 
      (track->Flag() <= mFlag[1]) ) {
    mGoodTrack = true;
  }

  if(track->TofBeta() > 0) {
    mIsTofTrack = true;
  }

  if(mSelType == 0) { //TPC selection
#ifndef TPC_DNDX    
      if( (mom >= mTpcMomentum[0]) && (mom <= mTpcMomentum[1]) &&
         ( track->NSigmaPion() <= mTpcNSigmaPion[0]) || (track->NSigmaPion() >= mTpcNSigmaPion[1]) &&
         ( (track->NSigmaElectron() <= mTpcNSigmaElectron[0]) || (track->NSigmaElectron() >= mTpcNSigmaElectron[1]) ) &&
         ( (track->NSigmaKaon() <= mTpcNSigmaKaon[0]) || (track->NSigmaKaon() >= mTpcNSigmaKaon[1]) ) &&
         ( (track->NSigmaProton() >= mTpcNSigmaProton[0]) && (track->NSigmaProton() <= mTpcNSigmaProton[1]) ) ) {

          mGoodPID = true;
      }
#else
      if( (track->P().mag() >= mTpcMomentum[0]) && (track->P().mag() <= mTpcMomentum[1]) &&
         (track->DndxNSigmaPion() <= mTpcNSigmaPion[0]) || (track->DndxNSigmaPion() >= mTpcNSigmaPion[1]) &&
         ( (track->DndxNSigmaElectron() <= mTpcNSigmaElectron[0]) || (track->DndxNSigmaElectron() >= mTpcNSigmaElectron[1]) ) &&
         ( (track->DndxNSigmaKaon() <= mTpcNSigmaKaon[0]) || (track->DndxNSigmaKaon() >= mTpcNSigmaKaon[1]) ) &&
         ( (track->DndxNSigmaProton() >= mTpcNSigmaProton[0]) && (track->DndxNSigmaProton() <= mTpcNSigmaProton[1]) ) ) {

          mGoodPID = true;
      }
#endif
  }
  else if(mSelType == 1) { //ToF selection

      if( mIsTofTrack && 
         (track->TofMassSqr() >= mTofMassSqr[0]) && (track->TofMassSqr() <= mTofMassSqr[1]) &&
         (track->P().mag() >= mTofMomentum[0]) && (track->P().mag() <= mTofMomentum[1]) ) {
          mGoodPID = true;
      }
  }
  else if(mSelType == 2) { //TnT selection

      if( mIsTofTrack &&
         (track->P().mag() >= mTnTMomentum[0]) && (track->P().mag() <= mTnTMomentum[1]) &&
         (track->TofMassSqr() >= mTofMassSqr[0]) && (track->TofMassSqr() <= mTofMassSqr[1]) &&
         (track->NSigmaProton() >= mTnTNSigmaProton[0]) && (track->NSigmaProton() <= mTnTNSigmaProton[1]) ) {
          mGoodPID = true;
      }
  }
  else if(mSelType == 3) { //if(ToF signal) {TnT selection} else {TPC selection}

      if(mIsTofTrack) {

#ifndef TPC_DNDX
          if( (track->P().mag() >= mTnTMomentum[0]) && (track->P().mag() <= mTnTMomentum[1]) &&
             (track->TofMassSqr() >= mTofMassSqr[0]) && (track->TofMassSqr() <= mTofMassSqr[1]) &&
             (track->NSigmaProton() >= mTnTNSigmaProton[0]) && (track->NSigmaProton() <= mTnTNSigmaProton[1]) ) {
              mGoodPID = true;
          }
#else
          if( (track->P().mag() >= mTnTMomentum[0]) && (track->P().mag() <= mTnTMomentum[1]) &&
             (track->TofMassSqr() >= mTofMassSqr[0]) && (track->TofMassSqr() <= mTofMassSqr[1]) &&
             (track->DndxNSigmaProton() >= mTnTNSigmaProton[0]) && (track->DndxNSigmaProton() <= mTnTNSigmaProton[1]) ) {
              mGoodPID = true;
          }
#endif
      }
      else {

#ifndef TPC_DNDX      
          if( (track->P().mag() >= mTpcMomentum[0]) && (track->P().mag() <= mTpcMomentum[1]) &&
             (track->NSigmaPion() <= mTpcNSigmaPion[0]) || (track->NSigmaPion() >= mTpcNSigmaPion[1]) &&
             ( (track->NSigmaElectron() <= mTpcNSigmaElectron[0]) || (track->NSigmaElectron() >= mTpcNSigmaElectron[1]) ) &&
             ( (track->NSigmaKaon() <= mTpcNSigmaKaon[0]) || (track->NSigmaKaon() >= mTpcNSigmaKaon[1]) ) &&
             ( (track->NSigmaProton() >= mTpcNSigmaProton[0]) && (track->NSigmaProton() <= mTpcNSigmaProton[1]) ) ) {
              mGoodPID = true;
          }
#else
          if( (track->P().mag() >= mTpcMomentum[0]) && (track->P().mag() <= mTpcMomentum[1]) &&
             (track->DndxNSigmaPion() <= mTpcNSigmaPion[0]) || (track->DndxNSigmaPion() >= mTpcNSigmaPion[1]) &&
             ( (track->DndxNSigmaElectron() <= mTpcNSigmaElectron[0]) || (track->DndxNSigmaElectron() >= mTpcNSigmaElectron[1]) ) &&
             ( (track->DndxNSigmaKaon() <= mTpcNSigmaKaon[0]) || (track->DndxNSigmaKaon() >= mTpcNSigmaKaon[1]) ) &&
             ( (track->DndxNSigmaProton() >= mTpcNSigmaProton[0]) && (track->DndxNSigmaProton() <= mTpcNSigmaProton[1]) ) ) {
              mGoodPID = true;
          }
#endif      
      }
  } //End with PID

#ifdef STHBTDEBUG
  std::cout << " gregsProtonTrackCut::Pass() [DEBUG] "
      << "mGoodCharge: " << mGoodCharge << " mGoodTrack: " << mGoodTrack
      << " mGoodPID: " << mGoodPID << std::endl;
#endif

  if(mGoodCharge && mGoodTrack && mGoodPID) {
      mPassTrack = true;
      mNTracksPassed++;
  }
  else {
      mNTracksFailed++;
  }

  return mPassTrack;
}

//_________________
StHbtString gregsProtonTrackCut::Report() {

    stringstream rep_str;
    rep_str << std::endl << "-- gregsProtonTrackCut --" << std::endl
        << "Particle mass: " << mMass << std::endl
        << " Charge: " << mCharge << std::endl
        << mP[0] << " <= p <= " << mP[1] << std::endl
        << mPt[0] << " <= pT <= " << mPt[1] << std::endl
        << mEta[0] << " <= eta <= " << mEta[1] << std::endl
        << mRapidity[0] << " <= y <= " << mRapidity[1] << std::endl
        << mNHits[0] << " <= nHits <= " << mNHits[1] << std::endl
        << mNHitsFit[0] << " <= nHitsFit <= " << mNHitsFit[1] << std::endl
        << "nHitsFit/nHitsPoss >= " << mAntiSplit << std::endl
        << "Selection type = " << mSelType << std::endl
        << "== TPC identification info:" << std::endl
        << mTpcMomentum[0] << " <= p <= " << mTpcMomentum[1] << std::endl
        << mTpcNSigmaPion[0] << " <= nSigma(pi) <= " << mTpcNSigmaPion[1] << std::endl
        << "nSigma(e) <= " << mTpcNSigmaElectron[0] << " ||  nSigma(e) >= " << mTpcNSigmaElectron[1] << std::endl
        << "nSigma(K) <= " << mTpcNSigmaKaon[0] << " ||  nSigma(K) >= " << mTpcNSigmaKaon[1] << std::endl
        << "nSigma(p) <= " << mTpcNSigmaProton[0] << " ||  nSigma(p) >= " << mTpcNSigmaProton[1] << std::endl
        << "== ToF identification info:" << std::endl
        << mTofMomentum[0] << " <= p <= " << mTofMomentum[1] << std::endl
        << mTofMassSqr[0] << " <= mTofMassSqr <= " << mTofMassSqr[1] << std::endl
        << mTofInvBetaDiff[0] << " <= mInvBetaDiff <= " << mTofInvBetaDiff[1] << std::endl
        << "== TnT identification info:" << std::endl
        << mTnTMomentum[0] << " <= p <= " << mTnTMomentum[1] << std::endl
        << mTnTNSigmaProton[0] << " <= nSigma(pi) <= " << mTnTNSigmaProton[1] << std::endl
        << " >>> N Tracks passed: " << mNTracksPassed << std::endl
        << " >>> N Tracks failed: " << mNTracksFailed << std::endl;

    return rep_str.str();
}