#include "StFemtoDstMaker_HFT.h"
#include "StBTofHeader.h"
#include "TBranch.h"
#include "StPhysicalHelixD.hh"
#include "StDcaGeometry.h"

ClassImp(StFemtoDstMaker_HFT)

//
// Set maximum file size to 1.9 GB (Root has a 2GB limit)
//
#define MAXFILESIZE 1900000000

//_________________
StFemtoDstMaker_HFT::StFemtoDstMaker_HFT(StMuDstMaker *muMaker,
                                         const Char_t *oFileName) {

  std::cout << "StFemtoDstMaker_HFT::StFemtoDstMaker_HFT - Creating an instance...";
  mOutFileName = oFileName;
  mMuDstMaker = muMaker;
  mRefMultCorrUtil = NULL;
  mCompression = 9;

  mCollisionType = false;     // 0-pp, 1-AuAu

  mEventIsGood = false;
  mNEventsIn = 0;
  mNEventsPassed = 0;

  mAuAuCorrectZdc = true;

  mIsGoodTrack = false;
  mCurrentTrigger = 0;

  //
  // Default cut values
  //

  // initialize event cut variables
  mPrimVtxZ[0] = -70.;
  mPrimVtxZ[1] = 70.;
  mPrimVtxR[0] = -1.;
  mPrimVtxR[1] = 10.;
  mPrimVtxVpdVzDiff[0] = -10.;
  mPrimVtxVpdVzDiff[1]= 10.;
  mPrimVtxXShift = 0.;
  mPrimVtxYShift = 0.;

  // initialize single-particle cut variables
  mTrackMomentum[0] = 0.1;
  mTrackMomentum[1] = 2.;
  mTrackDcaPerp[0] = 0.;
  mTrackDcaPerp[1] = 2.;
  mTrackDcaZ[0] = -3.;
  mTrackDcaZ[1] = 3.;
  mTrackNHits[0] = 15;
  mTrackNHits[1] = 50;
  mTrackEta[0] = -1.1;
  mTrackEta[1] = 1.1;
  mTrackFlag[0] = 0;
  mTrackFlag[1] = 1000;

  // flags and cuts for exclusion
  mRemovePions = false;
  mRemoveKaons = false;
  mRemoveProtons = false;
  mPthresh   = 0.4;
  mTpcPionNSigma[0] = -1.5;
  mTpcPionNSigma[1] = 1.5;
  mTofPionMassSqr[0] = -0.1;
  mTofPionMassSqr[1] = 0.15;
  mTpcKaonNSigma[0] = -1.5;
  mTpcKaonNSigma[1] = 1.5;
  mTofKaonMassSqr[0] = 0.2;
  mTofKaonMassSqr[1] = 0.35;
  mTpcProtonNSigma[0] = -1.5;
  mTpcProtonNSigma[1] = 1.5;
  mTofProtonMassSqr[0] = 0.7;
  mTofProtonMassSqr[1] = 1.1;

  std::cout << "\t[DONE]" << std::endl;
}

//_________________
StFemtoDstMaker_HFT::~StFemtoDstMaker_HFT() {
  /* nothing to do */
}

//_________________
Int_t StFemtoDstMaker_HFT::Init() {

  std::cout << "StFemtoDstMaker_HFT::Init - Initializing the maker"
            << std::endl 
            << "Creating output file: " << mOutFileName;

  mFemtoEvent = new StFemtoEvent();
  mOutFile = new TFile(mOutFileName, "RECREATE");
  mOutFile->SetCompressionLevel(mCompression);
  std::cout << "\t[DONE]" << std::endl;

  mTree = new TTree("StFemtoDst","StFemtoDst");
  mTree->SetMaxTreeSize(MAXFILESIZE);
  //mTree->SetAutoSave(1000000);
  mTree->Branch("StFemtoEvent","StFemtoEvent", &mFemtoEvent);

  if(mCollisionType == true) { // for AuAu only
    if(!mRefMultCorrUtil) {
      std::cout << "StFemtoDstMaker_HFT::Init - Initializing StRefMultCorr...";
      mRefMultCorrUtil = new StRefMultCorr("refmult");
      std::cout << "\t[DONE]" << endl;
    }
  }

  mNBytes = 0;
  std::cout << "StFemtoDstMaker_HFT::Init - Initialization has been finished"
      << std::endl;
  return StMaker::Init();
}

//________________
void StFemtoDstMaker_HFT::Clear(Option_t *option) {
  StMaker::Clear();
}

//________________
Int_t StFemtoDstMaker_HFT::Make() {

  mNEventsIn++;

  mMuDst = NULL;
  mMuEvent = NULL;

  //
  // MuDstMaker initialization
  //
  if(!mMuDstMaker) {
    mMuDstMaker = (StMuDstMaker*)GetMaker("MuDst");
    if (!mMuDstMaker) {
      LOG_ERROR << "StFemtoDstMaker_HFT::Make [ERROR] - Cannot find StMuDstMaker"
          << std::endl;
      return kStOk;
    }
  }

  //
  // Obtaining MuDst
  //
    mMuDst = mMuDstMaker->muDst();
    if (!mMuDst) {
      mMuDst = (StMuDst*)GetInputDS("MuDst");
      if (!mMuDst) {
        gMessMgr->Warning() << "StFemtoDstMaker_HFT::Make [WARNING] - No MuDst has been found" 
                            << endm;
        return kStOk;
      }
    }

  //
  // obtaining MuEvent
  //
  mMuEvent = (StMuEvent*)mMuDst->event();
  if(!AcceptTrigger(mMuEvent)) { // if trigger is not found
    return kStOk;
  }

  if(mMuEvent->refMult() < 0) { // multiplicity cannot be negative
    return kStOk;
  }

  Int_t mNVertices = mMuDst->numberOfPrimaryVertices();
  Int_t mNPrimTracks = mMuDst->numberOfPrimaryTracks();
  Int_t mNGlobTracks = mMuDst->numberOfGlobalTracks();

  //
  // Some initializations of local variables
  //
  mPrimVertex = NULL;
  StThreeVectorF mVertPosition;
  Float_t mVpdVz = 0.;
  StBTofHeader* mTofHeader = mMuDst->btofHeader();

  Int_t mCent16;
  Double_t mCentWeight;
  Double_t mRefMultCorrVal = 0.;
  Double_t mBeta;
  Double_t mMassSqr;
  Int_t    mPrimVertIndex = -999;
  Float_t  mRanking = -999.;

  if(mCollisionType == false) { // pp collisions means false state
    mCent16 = -1;
    mCentWeight = -999.;
  }

  CleanVariables(); // clean index vectors

  //
  // Vertex loop
  //
  for(Int_t iVert=0; iVert<mNVertices; iVert++) {

    mEventIsGood = false;

    if(iVert!=0) continue; // not first primary vertex does not contain tracks
                           //   with fast detectors (TOF)
    mPrimVertex = mMuDst->primaryVertex(iVert);

    if(mPrimVertex->ranking() <= 0) // positive ranking only
      continue;
    //
    // Check that position of the primary vertex if not (0,0,0)
    //
    if(mPrimVertex->position().x() == 0 &&
       mPrimVertex->position().y() == 0 &&
       mPrimVertex->position().z() == 0)      continue;

    if(mNPrimTracks < 0 || mNPrimTracks > 10000) // reasonable amount of tracks
      continue;

    mPrimVertIndex = iVert;
    mRanking = mPrimVertex->ranking();
    mVertPosition = mPrimVertex->position();

    if(mTofHeader)
      mVpdVz = mTofHeader->vpdVz();
    else
      mVpdVz = mVertPosition.z();

    if( !AcceptPrimVtx(mVertPosition, mVpdVz) )
      continue;

    mEventIsGood = true;
    //
    // Set StRefMultCorr
    //
    if(mCollisionType == true) { // AuAu collisions means true state
      mRefMultCorrUtil->init(mMuEvent->runId());
      if(mAuAuCorrectZdc == true) {
        mRefMultCorrUtil->initEvent(mMuEvent->refMult(),
                                    mVertPosition.z(),
                                    mMuEvent->runInfo().zdcCoincidenceRate());
      }
      else {
        mRefMultCorrUtil->initEvent(mMuEvent->refMult(),
                                    mVertPosition.z());
      }
      mCent16 = mRefMultCorrUtil->getCentralityBin16();
      if(mCent16 < 0) // some AuAu collisions have -1 value
        continue;
      mCentWeight = mRefMultCorrUtil->getWeight();
      mRefMultCorrVal = mRefMultCorrUtil->getRefMultCorr();
    } //if(mCollisionType == true)

    mNEventsPassed++;
    mFemtoEvent->SetEventId(mMuEvent->eventId());
    mFemtoEvent->SetRunId(mMuEvent->runId());
    mFemtoEvent->SetCollisionType(mCollisionType);
    mFemtoEvent->SetCent16(mCent16);
    //mFemtoEvent->SetZDCe(mMuEvent->zdcTriggerDetector().adc(4));
    //mFemtoEvent->SetZDCw(mMuEvent->zdcTriggerDetector().adc(0));
    mFemtoEvent->SetNumberOfBTofHit(mMuDst->numberOfBTofHit());
    mFemtoEvent->SetNumberOfPrimaryTracks(mNPrimTracks);
    mFemtoEvent->SetNumberOfGlobalTracks(mNGlobTracks);
    mFemtoEvent->SetMagneticField(mMuEvent->magneticField());
    mFemtoEvent->SetPrimaryVertexPosition(mVertPosition.x(),
                                          mVertPosition.y(),
                                          mVertPosition.z());
    mFemtoEvent->SetVpdVz(mVpdVz);
    mFemtoEvent->SetTriggerId(mCurrentTrigger);
    mFemtoEvent->SetPrimaryVertexRanking(mRanking);

    UInt_t refMult2 = 0;
    UInt_t refMult2Pos = 0;
    UShort_t refMult = 0;
    UShort_t refMultPos = 0;

    StFemtoTrack *mFTrack = NULL;

    //
    // Loop over primary tracks
    //
    for(Int_t iTrk=0; iTrk < mNGlobTracks; iTrk++) {

      mGlobTrack = mMuDst->globalTracks(iTrk);
      CalcDca(mGlobTrack->helix(), mPrimVertex->position());

      if (AcceptRefMult2(mGlobTrack)) {
        refMult2++;
        if(mGlobTrack->charge() > 0)    refMult2Pos++;
      }

      if (AcceptRefMult(mGlobTrack)) {
        refMult++;
        if(mGlobTrack->charge() > 0)    refMultPos++;
      }

      if(!AcceptTrack(mGlobTrack, iVert))
        continue;

      Bool_t mIsTofTrack = IsTofTrack(mGlobTrack);
      if(mIsTofTrack) {
        mBeta = mGlobTrack->btofPidTraits().beta();
        mMassSqr = mGlobTrack->momentum().mag2()*(1./(mBeta*mBeta) - 1.);
        //
        // Next lines are cheat. But we should remove garbage
        //
        if(mMassSqr < -0.02) continue;
        if(mGlobTrack->momentum().mag() > 0.5) {
          if(mMassSqr > 0.06 && mMassSqr < 0.18)   continue;
          if(mMassSqr > 0.35 && mMassSqr < 0.8)    continue;
          if(mMassSqr > 1.)                        continue;
        }
      } //if(mIsTofTrack)
      else {
        mBeta = -999.;
        mMassSqr = -999.;
      }

      Bool_t mIsTpcPion   = false;
      Bool_t mIsTpcKaon   = false;
      Bool_t mIsTpcProton = false;
      Bool_t mIsTofPion   = false;
      Bool_t mIsTofKaon   = false;
      Bool_t mIsTofProton = false;

      if(mRemovePions) {
        mIsTpcPion = IsTpcPion(mGlobTrack);
        mIsTofPion = IsTofPion(mGlobTrack);

        if(mIsTpcPion==true || mIsTofPion==true)
          continue;
      }

      if(mRemoveKaons) {
        mIsTpcKaon = IsTpcKaon(mGlobTrack);
        mIsTofKaon = IsTofKaon(mGlobTrack);

        if(mIsTpcKaon==true || mIsTofKaon==true)
          continue;
      }

      if(mRemoveProtons) {
        mIsTpcProton = IsTpcProton(mGlobTrack);
        mIsTofProton = IsTofProton(mGlobTrack);

        if(mIsTpcProton==true || mIsTofProton==true)
          continue;
      }

      mFTrack = mFemtoEvent->AddFemtoTrack();

      mFTrack->SetId( mGlobTrack->id() );
      mFTrack->SetNHits( (Char_t)(mGlobTrack->charge()*mGlobTrack->nHits()) );
      mFTrack->SetNHitsPoss( mGlobTrack->nHitsPoss() );
      mFTrack->SetNSigmaElectron( mGlobTrack->nSigmaElectron() );
      mFTrack->SetNSigmaPion( mGlobTrack->nSigmaPion() );
      mFTrack->SetNSigmaKaon( mGlobTrack->nSigmaKaon() );
      mFTrack->SetNSigmaProton( mGlobTrack->nSigmaProton() );
      mFTrack->SetDedx(mGlobTrack->dEdx() );
      mFTrack->SetMap0(mGlobTrack->topologyMap().data(0));
      mFTrack->SetMap1(mGlobTrack->topologyMap().data(1));
      mFTrack->SetP(0.0, 0.0, 0.0);
      mFTrack->SetDCAxGlobal(mDca.x());
      mFTrack->SetDCAyGlobal(mDca.y());
      mFTrack->SetDCAzGlobal(mDca.z());
      mFTrack->SetGlobalP(mGlobTrack->p().x(),
                          mGlobTrack->p().y(),
                          mGlobTrack->p().z());

      if(mIsTofTrack) {
        mFTrack->SetBeta(mBeta);
      }
      else {
        mFTrack->SetBeta(-999.);
      }
    }
    mFemtoEvent->SetRefMult(refMult);
    mFemtoEvent->SetRefMultCorr(mCollisionType ? mRefMultCorrVal : refMult);
    mFemtoEvent->SetRefMultCorrWeight(mCentWeight);
    mFemtoEvent->SetRefMultPos(refMultPos);
    mFemtoEvent->SetRefMult2(refMult2);
    mFemtoEvent->SetRefMult2Pos(refMult2Pos);

    if(mEventIsGood) {
      mNBytes += mTree->Fill();
      mFemtoEvent->Clear();
    }
  }

  return kStOk;
}

//_________________
Bool_t StFemtoDstMaker_HFT::AcceptTrigger(StMuEvent *muEvent) {
  //We will store the bit mask for the list of triggers.
  //The list of triggers should be the same and in the same
  //order, in order to read it later
  Bool_t mIsGoodTrigger = false;
  mCurrentTrigger = 0;
  if(mTriggerIdCollection.empty()) {
    mIsGoodTrigger = true;
  }
  else {
    for (UInt_t iTrg = 0; iTrg < mTriggerIdCollection.size(); iTrg++) {
      if(muEvent->triggerIdCollection().nominal().isTrigger(mTriggerIdCollection.at(iTrg))) {
        mIsGoodTrigger = true;
        mCurrentTrigger |= (1<<iTrg);
        break;
      }
    } //for(UInt_t iTrg=0; iTrg<mTriggerIdCollection.size(); iTrg++)
  }

  return mIsGoodTrigger;
}

//_________________
Bool_t StFemtoDstMaker_HFT::AcceptPrimVtx(StThreeVectorF vtxPos, 
                                                  Float_t vpdVz) {
  Float_t mVtxX = vtxPos.x() - mPrimVtxXShift;
  Float_t mVtxY = vtxPos.y() - mPrimVtxYShift;
  Float_t vtxR = TMath::Sqrt(mVtxX*mVtxX + mVtxY*mVtxY);
  Float_t vpdDiff = vtxPos.z() - vpdVz;

  Bool_t mIsGoodVtx =  vtxPos.z() >= mPrimVtxZ[0] &&
                       vtxPos.z() <= mPrimVtxZ[1] &&
                       vtxR >= mPrimVtxR[0] &&
                       vtxR <= mPrimVtxR[1] &&
                       vpdDiff >= mPrimVtxVpdVzDiff[0] &&
                       vpdDiff <= mPrimVtxVpdVzDiff[1];

  return mIsGoodVtx;
}

//_________________
Bool_t StFemtoDstMaker_HFT::AcceptTrack(StMuTrack *trk, UShort_t vtxInd) {

  Bool_t mIsGoodTrack = trk->vertexIndex() == vtxInd &&
                        trk->momentum().mag() >= mTrackMomentum[0] && trk->momentum().mag() <= mTrackMomentum[1] &&
                        trk->type() == 0 &&
                        trk->eta() >= mTrackEta[0] && trk->eta() <= mTrackEta[1] &&
                        trk->nHits() >= mTrackNHits[0] && trk->nHits() <= mTrackNHits[1] &&
                        trk->flag() >= mTrackFlag[0] && trk->flag() <= mTrackFlag[1] &&
                        mDca.perp() >= mTrackDcaPerp[0] &&
                        mDca.perp() <= mTrackDcaPerp[1] &&
                        mDca.z() >= mTrackDcaZ[0] &&
                        mDca.z() <= mTrackDcaZ[1];

  return mIsGoodTrack;
}

//_________________
Bool_t StFemtoDstMaker_HFT::AcceptRefMult2(StMuTrack *trk) {

  Bool_t mIsGoodTrack = trk->flag()> 0 &&
                        trk->charge() != 0 &&
                        trk->nHits() >= 10 &&
                        mDca.mag() < 3. &&
                        trk->momentum().mag() >= 1e-10 &&
                        TMath::Abs(trk->eta()) <= 1.0;

  return mIsGoodTrack;
}

//_________________
Bool_t StFemtoDstMaker_HFT::AcceptRefMult(StMuTrack *trk) {

  Bool_t mIsGoodTrack = trk->flag()> 0 && // trk->flag() <= 700 pile tracks
                        trk->charge() != 0 &&
                        trk->nHits() >= 10 &&
                        mDca.mag() < 3. &&
                        trk->momentum().mag() >= 1e-10 &&
                        TMath::Abs(trk->eta()) <= 0.5;

  return mIsGoodTrack;
}

//_________________
Bool_t StFemtoDstMaker_HFT::IsTofTrack(StMuTrack *trk) { //Only for globals

  Bool_t mIsTofHit = trk->btofPidTraits().beta() > 0 &&
                     trk->btofPidTraits().timeOfFlight() > 0;

  return mIsTofHit;
}

//_________________
Int_t StFemtoDstMaker_HFT::Finish() {

  mOutFile->cd();
  mOutFile->Write();
  //mTree->Write();
  mOutFile->Close();

  std::cout << "*************************************" << std::endl
      << "StFemtoDstMaker_HFT has been finished" << std::endl
      << "\t nEventsPassed   : " << mNEventsPassed  << std::endl
      << "\t nEventsProcessed: " << mNEventsIn << std::endl
      << "*************************************" << std::endl;

  return kStOk;
}

//_________________
void StFemtoDstMaker_HFT::CleanVariables() {

  mEventIsGood = false;
}

//_________________
Bool_t StFemtoDstMaker_HFT::IsTpcPion(StMuTrack *pTrk) {

  Bool_t isPion = false;

  if( pTrk->momentum().mag() <= mPthresh && 
     pTrk->nSigmaPion() >= mTpcPionNSigma[0] &&
     pTrk->nSigmaPion() <= mTpcPionNSigma[1] ) {
    isPion = true;
  }

  return isPion;
}

//_________________
Bool_t StFemtoDstMaker_HFT::IsTpcKaon(StMuTrack *pTrk) {

  Bool_t isKaon = false;

  if( pTrk->momentum().mag() <= mPthresh && 
     pTrk->nSigmaKaon() >= mTpcKaonNSigma[0] &&
     pTrk->nSigmaKaon() <= mTpcKaonNSigma[1] ) {
    isKaon = true;
  }

  return isKaon;
}

//_________________
Bool_t StFemtoDstMaker_HFT::IsTpcProton(StMuTrack *pTrk) {

  Bool_t isProton = false;

  if( pTrk->momentum().mag() <= mPthresh && 
     pTrk->nSigmaProton() >= mTpcProtonNSigma[0] &&
     pTrk->nSigmaProton() <= mTpcProtonNSigma[1] ) {
    isProton = true;
  }

  return isProton;
}

//_________________
Bool_t StFemtoDstMaker_HFT::IsTofPion(StMuTrack *gTrk) {

  Bool_t isPion = false;
  Float_t beta = -999.;
  Float_t mSqr = -999.;
  if( gTrk->btofPidTraits().beta() > 0 &&
     gTrk->momentum().mag() > mPthresh ) {
    beta = gTrk->btofPidTraits().beta();
    mSqr = gTrk->momentum().mag2() * (1./(beta*beta) - 1.);

    if(mSqr >= mTofPionMassSqr[0] &&
       mSqr <= mTofPionMassSqr[1]) {
      isPion = true;
    }
  }

  return isPion;
}

//_________________
Bool_t StFemtoDstMaker_HFT::IsTofKaon(StMuTrack *gTrk) {

  Bool_t isKaon = false;
  Float_t beta = -999.;
  Float_t mSqr = -999.;
  if( gTrk->btofPidTraits().beta() > 0 &&
     gTrk->momentum().mag() > mPthresh ) {
    beta = gTrk->btofPidTraits().beta();
    mSqr = gTrk->momentum().mag2() * (1./(beta*beta) - 1.);

    if(mSqr >= mTofKaonMassSqr[0] &&
       mSqr <= mTofKaonMassSqr[1]) {
      isKaon = true;
    }
  }

  return isKaon;
}

//_________________
StMuTrack *StFemtoDstMaker_HFT::FindPrimaryTrack(int idGlob, int nPrimTracks) {
  StMuTrack *track;
  for (int iTrk = 0; iTrk < nPrimTracks; iTrk++) {
    if ((track = mMuDst->primaryTracks(iTrk))->id() == idGlob)
      return track;
  }
  return 0;
}

//_________________
Bool_t StFemtoDstMaker_HFT::IsTofProton(StMuTrack *gTrk) {

  Bool_t isProton = false;
  Float_t beta = -999.;
  Float_t mSqr = -999.;
  if(gTrk->btofPidTraits().beta() > 0 &&
     gTrk->momentum().mag() > 0.5 ) {

    //gTrk->momentum().mag() > mPthresh ) {
    beta = gTrk->btofPidTraits().beta();
    mSqr = gTrk->momentum().mag2() * (1./(beta*beta) - 1.);

    if(mSqr >= mTofProtonMassSqr[0] &&
       mSqr <= mTofProtonMassSqr[1]) {
      isProton = true;
    }
  }

  return isProton;
}

//_________________
void StFemtoDstMaker_HFT::SetTriggerId(const UInt_t &id) {
  mTriggerIdCollection.push_back(id);
}

//_________________
//
// This method calculates the distance of the closes approach (DCA) from the
// one track to a 3D-point (B point, see below).
//
// For this methods we needs only two arguments
//  the helix of the track itself and the 3D-coord of the point (for example
//  primary vertex position).
//
void StFemtoDstMaker_HFT::CalcDca(const StPhysicalHelixD &helix,
                                  const StThreeVectorF   &vtxPos) {
  double pathLen = helix.pathLength(vtxPos, false); // this method returns the path length
                                                    //  along the helix to the point A
                                                    //  of the DCA. We need it further
  mDca = helix.at(pathLen) - vtxPos; // here we take the 3D-vector of the point A
                                     //  on the helix and subtract the B point,
                                     //  which is the coordinate of the, for example,
                                     //  primary vertex.
}
