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

#ifdef __ROOT__
ClassImp(gregsSimpleTrackCut)
#endif

//_________________
gregsSimpleTrackCut::gregsSimpleTrackCut() {

  mNTracksPassed = mNTracksFailed = 0;


  mTnTNSigmaElectron[0]  = -1e9;  mTnTNSigmaElectron[1]  = +1e9;
  mTnTNSigmaPion[0]  = -1e9;  mTnTNSigmaPion[1]  = +1e9;
  mTnTNSigmaKaon[0]  = -1e9;  mTnTNSigmaKaon[1]  = +1e9;
  mTnTNSigmaProton[0]= -1e9;  mTnTNSigmaProton[1]= +1e9;
  mEta[0]      = -1e9;   mEta[1]      = +1e9; 
  mRapidity[0] = -1e9;   mRapidity[1] = +1e9; 
  mP[0] = -1e9;  mP[1] = +1e9;
  mPt[0] = -1e9;  mPt[1] = +1e9;
  mPx[0] = -1e9;  mPx[1] = +1e9;
  mPy[0] = -1e9;  mPy[1] = +1e9;
  mPz[0] = -1e9;  mPz[1] = +1e9;
  mDCA[0] = -1e9;  mDCA[1] = +1e9;
  mFlag[0] = 0;   mFlag[1] = 1000;
  mDCAGlobal[0] = -1e9;  mDCAGlobal[1] = +1e9;
  mNHits[0] = 0; mNHits[1] = 60;
  mNHitsFit[0] = 0; mNHitsFit[1] = 60;
  mNdEdxHits[0] = 0; mNdEdxHits[1] = 60;
  mAntiSplit = 0.;
  mCharge = 0;
  mType = 1;           //0-global, 1-primary

  //TPC identification
  mTpcMomentum[0]=0.;  mTpcMomentum[1]=+1e9;
  mNSigmaElectron[0] = -1e9;  mNSigmaElectron[1] = +1e9;
  mNSigmaPion[0]     = -1e9;  mNSigmaPion[1]     = +1e9;
  mNSigmaKaon[0]     = -1e9;  mNSigmaKaon[1]     = +1e9;
  mNSigmaProton[0]   = -1e9;  mNSigmaProton[1]   = +1e9;

  //TOF identification
  mTofMassSqr[0] = -1e9; mTofMassSqr[1] = +1e9;
  mTofMomentum[0] = 0.; mTofMomentum[1] = +1e9;

  //TPC&&TOF
  mTnTMomentum[0]=0.; mTnTMomentum[1]=+1e9;
  mTnTNSigmaElectron[0]=-1e9; mTnTNSigmaElectron[1]=+1e9;
  mTnTNSigmaPion[0]=-1e9; mTnTNSigmaPion[1]=+1e9;
  mTnTNSigmaKaon[0]=-1e9; mTnTNSigmaKaon[1]=+1e9;
  mTnTNSigmaProton[0]=-1e9; mTnTNSigmaProton[1]=+1e9;

  //TPC&&TOF (P>Pthresh) || TPC (P<Pthresh)
  mTTTMomentumThresh = 0.4;

  mDetSelection = 3;   // 0-TPC, 1-ToF, 2-TPC+ToF, 3-ToF||TPC, 4-TPC&&TOF(P>Pthr)||TPC(P<Pthr)
}

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

  mNTracksPassed = mNTracksFailed = 0;
  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];
  mNdEdxHits[0] = c.mNdEdxHits[0];  mNdEdxHits[1] = c.mNdEdxHits[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];

  mDetSelection = c.mDetSelection;

  //TPC identification
  mTpcMomentum[0] = c.mTpcMomentum[0]; mTpcMomentum[1] = c.mTpcMomentum[1];
  mNSigmaElectron[0] = c.mNSigmaElectron[0];  mNSigmaElectron[1] = c.mNSigmaElectron[1];
  mNSigmaPion[0] = c.mNSigmaPion[0];  mNSigmaPion[1] = c.mNSigmaPion[1];
  mNSigmaKaon[0] = c.mNSigmaKaon[0];  mNSigmaKaon[1] = c.mNSigmaKaon[1];
  mNSigmaProton[0] = c.mNSigmaProton[0];  mNSigmaProton[1] = c.mNSigmaProton[1];

  //TOF identification
  mTofMomentum[0] = c.mTofMomentum[0]; mTofMomentum[1] = c.mTofMomentum[1];
  mTofMassSqr[0] = c.mTofMassSqr[0];  mTofMassSqr[1] = c.mTofMassSqr[1];

  //TPC&&TOF
  mTnTMomentum[0] = c.mTnTMomentum[0]; mTnTMomentum[1] = c.mTnTMomentum[1];
  mTnTNSigmaElectron[0] = c.mTnTNSigmaElectron[0]; mTnTNSigmaElectron[1] = c.mTnTNSigmaElectron[1];
  mTnTNSigmaPion[0] = c.mTnTNSigmaPion[0]; mTnTNSigmaPion[1] = c.mTnTNSigmaPion[1];
  mTnTNSigmaKaon[0] = c.mTnTNSigmaKaon[0]; mTnTNSigmaKaon[1] = c.mTnTNSigmaKaon[1];
  mTnTNSigmaProton[0] = c.mTnTNSigmaProton[0]; mTnTNSigmaProton[1] = c.mTnTNSigmaProton[1];

  //TPC&&TOF (P>Pthresh) || TPC (P<Pthresh)
  mTTTMomentumThresh = c.mTTTMomentumThresh;
}

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

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

  //mMass has to be set in the macros and defined in the base class
  float tEnergy = ::sqrt(track->P().mag2() + mMass*mMass);
  float tRapidity = 0.5*::log( (tEnergy+track->P().z())/(tEnergy-track->P().z()) );
  float tPseudoRap = 0.5*::log( (track->P().mag()+track->P().z())/(track->P().mag()-track->P().z()) );
  float tAntiSplit = (float)track->NHitsFit()/track->NHitsPossible();

  //Main patterns
  bool mGoodCharge = false;
  bool mGoodTrack = false;
  bool mGoodTpcPID = false;
  bool mGoodTofPID = false;
  bool mGoodToTPID = false;
  bool mGoodTnTPID = false;
  bool mGoodTTTPID = false;
  bool mGoodPID = false;
  bool mPassTrack = false;

  if(mDetSelection<0 || mDetSelection>4) {
    std::cout << "gregsSimpleTrackCut::Pass[WARNING] Wrong detector selection type: " 
	      << mDetSelection << " . Resetting to 3" << std::endl;
  }

  //If is right charge
  mGoodCharge = (track->Charge() == mCharge);

  //If is normal track
  mGoodTrack = ( (track->TrackType() == mType) &&
		 (tAntiSplit >= 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]) &&
		 (track->NHitsDedx() >= mNdEdxHits[0]) && (track->NHitsDedx() <= mNdEdxHits[1]) &&
		 (track->P().mag() >= mP[0]) && (track->P().mag() <= 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]) &&
		 (tPseudoRap >= mEta[0]) && (tPseudoRap <= mEta[1]) &&
		 (tRapidity >= mRapidity[0]) && (tRapidity <= mRapidity[1]) &&
		 (track->Flag() >= mFlag[0]) && (track->Flag() <= mFlag[1])
		 );

  //If is good TPC identification: TPC nSigma always smaller the nTofNSigma
  mGoodTpcPID = ( (track->NSigmaElectron() >= mNSigmaElectron[0]) &&
		  (track->NSigmaElectron() <= mNSigmaElectron[1]) &&
		  (track->NSigmaPion() >= mNSigmaPion[0]) &&
		  (track->NSigmaPion() <= mNSigmaPion[1]) &&
		  (track->NSigmaKaon() >= mNSigmaKaon[0]) &&
		  (track->NSigmaKaon() <= mNSigmaKaon[1]) &&
		  (track->NSigmaProton() >= mNSigmaProton[0]) &&
		  (track->NSigmaProton() <= mNSigmaProton[1]) &&
		  (track->P().mag() >= mTpcMomentum[0]) &&
		  (track->P().mag() <= mTpcMomentum[1])
		  );

  //If is good TOF identification (in case when the TOF hit does exist)
  if(track->TofBeta() > 0) {
    mGoodTofPID = ( (track->TofMassSqr() >= mTofMassSqr[0]) &&
		    (track->TofMassSqr() <= mTofMassSqr[1]) &&
		    (track->P().mag() >= mTofMomentum[0]) &&
		    (track->P().mag() <= mTofMomentum[1])
		    );
  }

  //TOF identification (if ToF hit), and TPC (if no ToF hit)
  if(track->TofBeta() > 0) {
    mGoodToTPID = mGoodTofPID;
  }
  else {
    mGoodToTPID = mGoodTpcPID;
  }

  //TPC&&TOF identification
  mGoodTnTPID = ( (track->TofBeta() > 0) && 
		  (track->NSigmaElectron() >= mTnTNSigmaElectron[0]) &&
		  (track->NSigmaElectron() <= mTnTNSigmaElectron[1]) &&
		  (track->NSigmaPion() >= mTnTNSigmaPion[0]) &&
		  (track->NSigmaPion() <= mTnTNSigmaPion[1]) &&
		  (track->NSigmaKaon() >= mTnTNSigmaKaon[0]) &&
		  (track->NSigmaKaon() <= mTnTNSigmaKaon[1]) &&
		  (track->NSigmaProton() >= mTnTNSigmaProton[0]) &&
		  (track->NSigmaProton() <= mTnTNSigmaProton[1]) &&
		  (track->TofMassSqr() >= mTofMassSqr[0]) &&
		  (track->TofMassSqr() <= mTofMassSqr[1]) &&
		  (track->P().mag() >= mTnTMomentum[0]) &&
		  (track->P().mag() <= mTnTMomentum[1])
		  );

  //TPC&&TOF (P>Pthresh) || TPC (P<Pthresh)
  if( (track->P().mag() >= mTTTMomentumThresh) &&
      (track->P().mag() <= mTofMomentum[1])  ) {
    if(track->TofBeta() > 0) {
	mGoodTTTPID = mGoodTnTPID;
    }
  }
  else if( (track->P().mag() >= mTpcMomentum[0]) &&
	   (track->P().mag() < mTTTMomentumThresh) ){
    mGoodTTTPID = mGoodTpcPID;
  }


  //Detector selection criteria
  switch(mDetSelection) {
  case 0: //TPC only
    if(mGoodTpcPID == 1) {
      mGoodPID = true;
    }
    break;
  case 1: //ToF only
    if(mGoodTofPID == 1) {
      mGoodPID = true;
    }
    break;
  case 2: //TPC && ToF
    if(mGoodTnTPID == 1) {
      mGoodPID = true;
    }
    break;
  case 3: //ToF, if no ToF signal then use TPC
    if(mGoodToTPID == 1) {
      mGoodPID = true;
    }
    break;
  case 4:
    if(mGoodTTTPID) {
      mGoodPID = true;
    }
    break;
  default:
    if(mGoodToTPID) {
      mGoodPID = true;
    }
    break;
  };

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

  return mPassTrack;
}

//_________________
StHbtString gregsSimpleTrackCut::Report() {

  stringstream rep_str;
  rep_str << std::endl << "-- gregsSimpleTrackCut --" << std::endl
	  << "Charge: " << mCharge << std::endl
	  << "Selection type = " << mDetSelection << std::endl
	  << "TPC identification:" << std::endl
	  << mTpcMomentum[0] << " <= mTpcMomentum <= " << mTpcMomentum[1] << std::endl
	  << mNSigmaElectron[0] << " <= nSigma(e) <= " << mNSigmaElectron[1] << std::endl
	  << mNSigmaPion[0] << " <= nSigma(pi) <= " << mNSigmaPion[1] << std::endl
	  << mNSigmaKaon[0] << " <= nSigma(K) <= " << mNSigmaKaon[1] << std::endl
	  << mNSigmaProton[0] << " <= nSigma(p) <= " << mNSigmaProton[1] << std::endl
	  << "TOF identification: " << std::endl
	  << mTofMomentum[0] << " <= mTofMomentum <= " << mTofMomentum[1] << std::endl
	  << mTofMassSqr[0] << " <= mTofMassSqr <= " << mTofMassSqr[1] << std::endl
	  << "TnT identification: " << std::endl
	  << mTnTMomentum[0] << " <= mTnTMomentum <= " << mTnTMomentum[1] << std::endl
	  << mTnTNSigmaElectron[0] << " <= nTnTSigma(e) <= " << mTnTNSigmaElectron[1] << std::endl
	  << mTnTNSigmaPion[0] << " <= nTnTSigma(pi) <= " << mTnTNSigmaPion[1] << std::endl
	  << mTnTNSigmaKaon[0] << " <= nTnTSigma(K) <= " << mTnTNSigmaKaon[1] << std::endl
	  << mTnTNSigmaProton[0] << " <= nTnTSigma(p) <= " << mTnTNSigmaProton[1] << std::endl
	  << "TPC&TOF || TPC: " << std::endl
	  << "Pthreshold = " << mTTTMomentumThresh << std::endl
	  << "Genereal cut values:" << std::endl
	  << mNHits[0] << " <= NHits <= " << mNHits[1] << std::endl
	  << mNHitsFit[0] << " <= NHitsFit <= " << mNHitsFit[1] << std::endl
	  << mNdEdxHits[0] << " <= NHits <= " << mNdEdxHits[1] << std::endl
	  << mPt[0] << " <= pT <= " << mPt[1] << std::endl
	  << mP[0] << " <= p <= " << mP[1] << std::endl
	  << mRapidity[0] << " <= Rapidity <= " << mRapidity[1] << std::endl
	  << mEta[0] << " <= PseudoRapidity <= " << mEta[1] << std::endl
	  << "nHitsFit/nHitsPoss >= " << mAntiSplit << std::endl
	  << " >>> N Tracks passed: " << mNTracksPassed << std::endl
	  << " >>> N Tracks failed: " << mNTracksFailed << std::endl;

  return ( rep_str.str() );
}
