#include "StFemtoDstPicoRun11Maker.h"



ClassImp(StFemtoDstPicoRun11Maker)



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

  
StFemtoDstPicoRun11Maker::StFemtoDstPicoRun11Maker(StPicoDstMaker *picoMaker,
                                                   const Char_t *oFileName)
{
  std::cout << "StFemtoDstPicoRun11Maker::StFemtoDstPicoRun11Maker - Creating an instance...";
  mOutFileName = oFileName;
  //
  // PicoDstMaker initialization
  //
  if(!picoMaker)
    LOG_ERROR << "StFemtoDstPicoRun11Maker::Constructor[ERROR] - Cannot find StPicoMaker"
              << std::endl;
  mPicoMaker = picoMaker;
  mCompression = 9;

  mNEventsIn = 0;
  mNEventsPassed = 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.;
  mTrackDcaGlobal[0] = 0.;
  mTrackDcaGlobal[1] = 5.;
  mTrackNHitsFit[0] = 15;
  mTrackNHitsFit[1] = 50;
  mTrackEta[0] = -1.1;
  mTrackEta[1] = 1.1;

  // 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;

  mRandom = new TRandom3(0);

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



StFemtoDstPicoRun11Maker::~StFemtoDstPicoRun11Maker()
{
  /* NO-OP */
}



int StFemtoDstPicoRun11Maker::Init()
{
  std::cout << "StFemtoDstPicoRun11Maker::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->Branch("StFemtoEvent","StFemtoEvent", &mFemtoEvent);

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



void StFemtoDstPicoRun11Maker::Clear(Option_t *option)
{
  StMaker::Clear();
}



int StFemtoDstPicoRun11Maker::Make()
{
  int refMult;
  unsigned int nTracks, nPrimTracks = 0;
  float vpdVz, ranking;
  unsigned int refMult2 = 0;
  unsigned int refMult2Pos = 0;
  StFemtoTrack *mFTrack;

  mNEventsIn++;

  //
  // Obtaining StPicoDst
  //
  StPicoDst *picoDst = mPicoMaker->picoDst();
  if (!picoDst)
  {
    std::cout << "can't obtain StPicoDst" << std::endl;
    return kStWarn;
  }

  //
  // Obtaining PicoEvent
  //
  StPicoEvent *picoEvent = picoDst->event();
  if (!picoEvent)
  {
    std::cout << "can't obtain StPicoEvent" << std::endl;
    return kStWarn;
  }

  //
  // Check trigger
  //
  unsigned int trig = (unsigned int)picoEvent->triggerWord();
  /* For now it's not needed
  if (trig & (!mTrigWord))   return kStOk;
  */

  refMult = picoEvent->refMult();
  if (refMult < 0) // multiplicity cannot be negative
    return kStOk;

  nTracks = picoDst->numberOfTracks();
  ranking = picoEvent->ranking();
  if(ranking <= 0) // positive ranking only
    return kStOk;

  //
  // Check that position of the primary vertex if not (0,0,0)
  //
  const StThreeVectorF &primVtxPos = picoEvent->primaryVertex();
  vpdVz = picoEvent->vzVpd();
  if (primVtxPos.x() == 0 &&
      primVtxPos.y() == 0 &&
      primVtxPos.z() == 0)   return kStOk;
  if (!AcceptPrimVtx(primVtxPos, vpdVz))   return kStOk;

  mNEventsPassed++;
  mFemtoEvent->SetEventId(picoEvent->eventId());
  mFemtoEvent->SetRunId(picoEvent->runId());
  mFemtoEvent->SetCollisionType(false); // false means p+p
  mFemtoEvent->SetRefMult(refMult);
  mFemtoEvent->SetRefMultCorr((double)refMult);
  mFemtoEvent->SetRefMultCorrWeight(-999.);
  mFemtoEvent->SetRefMultPos(picoEvent->refMultPos());
  mFemtoEvent->SetCent16(-1);
  mFemtoEvent->SetNumberOfBTofHit(picoEvent->btofTrayMultiplicity());
  mFemtoEvent->SetNumberOfPrimaryTracks(nPrimTracks);
  mFemtoEvent->SetNumberOfGlobalTracks(picoEvent->numberOfGlobalTracks());
  mFemtoEvent->SetMagneticField(picoEvent->bField());
  mFemtoEvent->SetPrimaryVertexPosition(primVtxPos.x(),
                                        primVtxPos.y(),
                                        primVtxPos.z());
  mFemtoEvent->SetVpdVz(vpdVz);
  mFemtoEvent->SetTriggerId(trig);
  mFemtoEvent->SetPrimaryVertexRanking(ranking);

  //
  // Loop over primary tracks
  //
  for (unsigned int iTrk = 0; iTrk < nTracks; iTrk++)
  {
    mPicoTrack = picoDst->track(iTrk);
    mPrimP = mPicoTrack->pMom();
    mGlobP = mPicoTrack->gMom();
    short charge = mPicoTrack->charge();
    float beta = mPicoTrack->btofBeta();

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

    if(!AcceptTrack())   continue;

    //
    // Calculate mass square
    //
    mIsTofTrack = beta >= 1e-10;
    if(mIsTofTrack)
    {
      mMassSqr = mPrimP.mag2()*(1./(beta*beta) - 1.);
      //
      // Next lines are cheat. But we should remove garbage
      //
      if(mMassSqr < -0.02)   continue;
      if(mPrimP.mag() > 0.5) // in this range TOF efficienty ~>98%
      {
        if(mMassSqr > 0.06 && mMassSqr < 0.18)   continue;
        if(mMassSqr > 0.35 && mMassSqr < 0.8)    continue;
        if(mMassSqr > 1.)                        continue;
      }
    }
    else
    {
      beta = -999.;
      mMassSqr = -999.;
    }

    //
    // Remove unwanted particles
    //
    if(mRemovePions)
      if (IsTpcPion() || IsTofPion())   continue;

    if(mRemoveKaons)
      if (IsTpcKaon() || IsTofKaon())   continue;

    if(mRemoveProtons)
      if (IsTpcProton() || IsTofProton())   continue;


    //
    // Add new femto track and fill it with data
    //
    mFTrack = mFemtoEvent->AddFemtoTrack();

    mFTrack->SetId(mPicoTrack->id());
    mFTrack->SetNHits((Char_t)(charge*mPicoTrack->nHitsFit()));
    mFTrack->SetNHitsPoss(mPicoTrack->nHitsMax());
    mFTrack->SetNSigmaElectron(mPicoTrack->nSigmaElectron());
    mFTrack->SetNSigmaPion(mPicoTrack->nSigmaPion());
    mFTrack->SetNSigmaKaon(mPicoTrack->nSigmaKaon());
    mFTrack->SetNSigmaProton(mPicoTrack->nSigmaProton());
    mFTrack->SetDedx(mPicoTrack->dEdx()*1e-6);

    //
    // Calculate topology map
    //
    TopologyMapCalc();
    mFTrack->SetMap0(mTmap0);
    mFTrack->SetMap1(mTmap1);
    mFTrack->SetP(mPrimP.x(), mPrimP.y(), mPrimP.z());

    //
    // Calculate DCA
    //
    StPhysicalHelixD globHelix(mGlobP,
                               mPicoTrack->origin(),
                               picoEvent->bField()*kilogauss,
                               charge);
    CalcDca(globHelix, primVtxPos);

    mFTrack->SetDCAxGlobal(mDca.x());
    mFTrack->SetDCAyGlobal(mDca.y());
    mFTrack->SetDCAzGlobal(mDca.z());

    mFTrack->SetGlobalP(mGlobP.x(), mGlobP.y(), mGlobP.z());
    mFTrack->SetBeta(beta);
  }

  mFemtoEvent->SetRefMult2(refMult2);
  mFemtoEvent->SetRefMult2Pos(refMult2Pos);

  mNBytes += mTree->Fill();
  mFemtoEvent->Clear();

  return kStOk;
}



bool StFemtoDstPicoRun11Maker::AcceptPrimVtx(const StThreeVectorF &vtxPos,
                                             float vpdVz)
{
  float vtxR = vtxPos.perp();
  float vpdDiff = vtxPos.z() - vpdVz;

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



bool StFemtoDstPicoRun11Maker::AcceptTrack()
{
  return mPrimP.mag() >= mTrackMomentum[0] && mPrimP.mag() <= mTrackMomentum[1] &&
         mPrimP.pseudoRapidity() >= mTrackEta[0] && mPrimP.pseudoRapidity() <= mTrackEta[1] &&
         mPicoTrack->nHitsFit() >= mTrackNHitsFit[0] && mPicoTrack->nHitsFit() <= mTrackNHitsFit[1] &&
         mPicoTrack->dca() >= mTrackDcaGlobal[0] && mPicoTrack->dca() <= mTrackDcaGlobal[1];
}



bool StFemtoDstPicoRun11Maker::AcceptRefMult2()
{
  return mPicoTrack->charge() != 0 &&
         mPicoTrack->nHitsFit() >= 10 &&
         mPicoTrack->dca() < 3.0 &&
         mPrimP.mag() >= 1e-10 &&
         TMath::Abs(mPrimP.pseudoRapidity()) <= 1.0;
}



int StFemtoDstPicoRun11Maker::Finish()
{

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

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

  return kStOk;
}



void StFemtoDstPicoRun11Maker::TopologyMapCalc()
{
  if (mPrimP.mag() < 0.15) // tracks only in the inner sector. R = 1 m
  {
    mTmap0 = 0;
    mTmap1 = 0;
    for (int iBit = 0; iBit < 32; iBit++)
    {
      if (mRandom->Rndm() < 0.95)   mTmap0 |= 1 << iBit;
      else                          mTmap0 |= 0 << iBit;
    }

    mTmap0 |= 1 << 0; // default set to the first bit: primary-vertex-used
  }
  else
  {
    if (mPrimP.mag() < 0.3) // tracks may come to the end of the outer sector: R = 2 m
    {
      mTmap0 = 0;
      mTmap1 = 0;
      for (int iBit=0; iBit<32; iBit++)
      {
        if (mRandom->Rndm() < 0.95)   mTmap0 |= 1 << iBit;
        else                          mTmap0 |= 0 << iBit;

        if (iBit < 27)
        {
          if (mRandom->Rndm() < 0.85)   mTmap1 |= 1 << iBit;
          else                          mTmap1 |= 0 << iBit;
        }
      }
      mTmap0 |= 1 << 0; // default set to the first bit: primary-vertex-used
    }
    else
    {
      mTmap0 = 0;
      mTmap1 = 0;
      for (int iBit=0; iBit<32; iBit++)
      {
        if (mRandom->Rndm() < 0.95)   mTmap0 |= 1 << iBit;
        else                          mTmap0 |= 0 << iBit;

        if (iBit < 27)
        {
          if (mRandom->Rndm() < 0.95)   mTmap1 |= 1 << iBit;
          else                          mTmap1 |= 0 << iBit;
        }
      }
      mTmap0 |= 1 << 0; // default set to the first bit: primary-vertex-used
    }
  }
}



bool StFemtoDstPicoRun11Maker::IsTpcPion()
{
  return mPrimP.mag() <= mPthresh &&
         mPicoTrack->nSigmaPion() >= mTpcPionNSigma[0] &&
         mPicoTrack->nSigmaPion() <= mTpcPionNSigma[1];
}



bool StFemtoDstPicoRun11Maker::IsTpcKaon()
{
  return mPrimP.mag() <= mPthresh &&
         mPicoTrack->nSigmaKaon() >= mTpcKaonNSigma[0] &&
         mPicoTrack->nSigmaKaon() <= mTpcKaonNSigma[1];
}



bool StFemtoDstPicoRun11Maker::IsTpcProton()
{
  return mPrimP.mag() <= mPthresh &&
         mPicoTrack->nSigmaProton() >= mTpcProtonNSigma[0] &&
         mPicoTrack->nSigmaProton() <= mTpcProtonNSigma[1];
}



bool StFemtoDstPicoRun11Maker::IsTofPion()
{
  return mIsTofTrack && mPrimP.mag() > mPthresh &&
         mMassSqr >= mTofPionMassSqr[0] &&
         mMassSqr <= mTofPionMassSqr[1];
}



bool StFemtoDstPicoRun11Maker::IsTofKaon()
{
  return mIsTofTrack && mPrimP.mag() > mPthresh &&
         mMassSqr >= mTofKaonMassSqr[0] &&
         mMassSqr <= mTofKaonMassSqr[1];
}



bool StFemtoDstPicoRun11Maker::IsTofProton()
{
  return mIsTofTrack && mPrimP.mag() > mPthresh &&
         mMassSqr >= mTofProtonMassSqr[0] &&
         mMassSqr <= mTofProtonMassSqr[1];
}



void StFemtoDstPicoRun11Maker::CalcDca(const StPhysicalHelixD &helix,
                                       const StThreeVectorF   &vtxPos)
{
  double pathLen = helix.pathLength(vtxPos, false);
  mDca = helix.at(pathLen) - vtxPos;
}
