#include "StFlow.h"
#include "StBTofHeader.h"
#include "TBranch.h"

ClassImp(StFlow)

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



//_________________
StFlow::StFlow(StMuDstMaker *muMaker, const Char_t *oFileName) 
{
  std::cout << "StFlow::StFlow - Creating an instance...";
  mMuDstMaker = muMaker;

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

  mIsGoodTrack = false;
  mCurrentTrigger = 0;

  mFileName = oFileName;
  //
  // Initialize event cut variables
  //
  mPrimVtxZ[0] = -30.;
  mPrimVtxZ[1] = 30.;
  mPrimVtxR[0] = 0.;
  mPrimVtxR[1] = 2.;
  mPrimVtxVpdVzDiff[0] = -5.;
  mPrimVtxVpdVzDiff[1]= 5.;
  mPrimVtxXShift = 0.;
  mPrimVtxYShift = 0.;
  //
  // Initialize single-particle cut variables
  //
  mTrackP[0] = 0.1;
  mTrackP[1] = 2.;
  mTrackDca[0] = 0.;
  mTrackDca[1] = 2.;
  mTrackDcaGlobal[0] = 0.;
  mTrackDcaGlobal[1] = 2.;
  mTrackNHits[0] = 15;
  mTrackNHits[1] = 50;
  mTrackNHitsFit[0] = 15;
  mTrackNHitsFit[1] = 50;

  mTrackEta[0] = -1.;
  mTrackEta[1] = 1.;
  mTrackFlag[0] = 0;
  mTrackFlag[1] = 1000;

  mRefMultCorrUtil = new StRefMultCorr("refmult");

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

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

//_________________
Int_t StFlow::Init() 
{
  std::cout << "StFlow::Init - Initializing the maker" << std::endl;

  //
  // Create histograms and output file
  //
  mOutFile = new TFile(mFileName, "RECREATE");
  
  hdV2vsCent = new TProfile("hdV2vsCent", ";Centrality (%);v_{2}",
                            10, 0., 100.);
  hV2vsCent = new TH2F("hV2vsCent", ";Centrality (%);v_{2}",
                       10, 0., 100.,
                       100, -1., 1.);

  std::cout << "StFlow::Init - Initialization has been finished" << std::endl;
  return StMaker::Init();
}

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

//________________
Int_t StFlow::Make() 
{
  mNEventsIn++;

  mMuDst = NULL;
  mMuEvent = NULL;

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

  //Obtaining MuDst
  mMuDst = (StMuDst*)GetInputDS("MuDst");
  if(!mMuDst) 
  {
    gMessMgr->Warning() << "StFlow::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;
  }

  //Multiplicity cannot be negative
  unsigned short refMult = mMuEvent->refMult();
  unsigned short refMultPos = 0;
  unsigned short refMultNeg = 0;

  if(refMult < 0) 
  {
    return kStOk;
  }

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

  //Some initializations of local variables
  mPrimVertex = NULL;
  StThreeVectorF mVertPosition;
  Float_t mVpdVz = 0.;
  Int_t   mPrimVertIndex = -999;
  Float_t mRanking = -999.;

  //Clean index vectors
  CleanVariables();

  //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);
    //Positive ranking only
    if(mPrimVertex->ranking() <= 0) continue;
    //Not (0,0,0) position of the primary vertex
    if(mPrimVertex->position().x() == 0 &&
       mPrimVertex->position().y() == 0 &&
       mPrimVertex->position().z() == 0) continue;
    //Reasonable amount of tracks
    if(mNPrimTracks < 0 || mNPrimTracks > 10000) continue;

    mPrimVertIndex = iVert;
    mVertPosition = mPrimVertex->position();
    mVpdVz = mMuDst->btofHeader()->vpdVz();

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

    Int_t mCent9;
    Int_t mCent16;
    mRefMultCorrUtil->init(mMuEvent->runId());
    mRefMultCorrUtil->initEvent(mMuEvent->refMult(),
                                mVertPosition.z(),
                                mMuEvent->runInfo().zdcCoincidenceRate());
    mCent9 = mRefMultCorrUtil->getCentralityBin9();
    if(mCent9<0) continue;   //Some AuAu collisions have -1 value

    float centBin = 0;
    if (mCent9 == 8 || mCent9 == 7)
      centBin = 0;
    else
      centBin = 8 - (mCent9 + 1);
    centBin = 5 + centBin*10.;
    mEventIsGood = true;

    //
    //Loop over primary tracks
    //
    float v2 = 0.;
    unsigned long steps = 0;
    for (int iTrk1 = 0; iTrk1 < mNPrimTracks - 1; iTrk1++) 
    {
      mPrimTrack = mMuDst->primaryTracks(iTrk1);
      mGlobTrack = mMuDst->globalTracks(mPrimTrack->index2Global());
      float phi1 = mPrimTrack->phi();

      if( !AcceptTrack(mPrimTrack, iVert) ) continue;
      for (int iTrk2 = iTrk1 + 1; iTrk2 < mNPrimTracks; iTrk2++)
      {
        mPrimTrack = mMuDst->primaryTracks(iTrk2);
        if( !AcceptTrack(mPrimTrack, iVert) ) continue;

        float phi2 = mPrimTrack->phi();
        float dv2 = cos(2*(phi1 - phi2));
        v2 += dv2;
        steps++;
        hdV2vsCent->Fill(centBin, dv2);
      }
    } //for(Int_t iTrk=0; iTrk<mNPrimTracks; iTrk++)
    v2 /= (float)steps;
    hV2vsCent->Fill(centBin, v2);
  } //for(Int_t iVert=0; iVert<mNVertices; iVert++)

  return kStOk;
}

//_________________
Bool_t StFlow::AcceptTrigger(StMuEvent *muEvent) 
{
  Bool_t mIsGoodTrigger = false;
  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 = mTriggerIdCollection.at(iTrg);
        break;
      }
    } //for(UInt_t iTrg=0; iTrg<mTriggerIdCollection.size(); iTrg++)
  }

  return mIsGoodTrigger;
}

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

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

//_________________
Bool_t StFlow::AcceptTrack(StMuTrack *trk, UShort_t vtxInd) 
{
  float          mom      = trk->momentum().mag();
  float          eta      = trk->eta();
  unsigned short nHitsFit = trk->nHitsFit();
  float          ratioH   = (float)trk->nHitsFit()/(float)trk->nHitsPoss();
  short          flag     = trk->flag();
  float          dca      = trk->dcaGlobal(vtxInd).perp();

  bool isMomOk      = ( mom >= mTrackP[0] &&
                        mom <= mTrackP[1] );
  bool isEtaOk      = ( eta >= mTrackEta[0] &&
                        eta <= mTrackEta[1] );
  bool isNHitsFitOk = ( nHitsFit >= mTrackNHitsFit[0] &&
                        nHitsFit <= mTrackNHitsFit[1] &&
                        ratioH >= 0.51 );
  bool isFlagOk     = ( flag >= mTrackFlag[0] &&
                        flag <= mTrackFlag[1] );
  bool isDcaOk      = ( dca >= mTrackDcaGlobal[0] &&
                        dca <= mTrackDcaGlobal[1] );

  return ( isMomOk && isEtaOk && isNHitsFitOk && isFlagOk && isDcaOk );
}


//_________________
Int_t StFlow::Finish() 
{

  std::cout << "*************************************" << std::endl
      << "StFlow has been finished" << std::endl
      << "\t nEventsPassed   : " << mNEventsPassed  << std::endl
      << "\t nEventsProcessed: " << mNEventsIn << std::endl
      << "*************************************" << std::endl;
  mOutFile->cd();
  mOutFile->Write();
  mOutFile->Close();

  return kStOk;
}

//_________________
void StFlow::CleanVariables() 
{

  mEventIsGood = false;
}

//_________________
void StFlow::SetTriggerId(unsigned int id) 
{
  mTriggerIdCollection.push_back(id);
}


//_________________
void StFlow::SetMuDstMaker(StMuDstMaker *maker) 
{
  mMuDstMaker = maker;
}


//_________________
void StFlow::SetVtxZCut(float lo, float hi) 
{
  mPrimVtxZ[0] = lo;
  mPrimVtxZ[1] = hi;
}


//_________________
void StFlow::SetVtxRCut(float lo, float hi) 
{
  mPrimVtxR[0] = lo;
  mPrimVtxR[1] = hi;  
}


//_________________
void StFlow::SetVtxShift(float xShift, float yShift) 
{
  mPrimVtxXShift = xShift;
  mPrimVtxYShift = yShift;
}


//_________________
void StFlow::SetVtxVpdVzDiffCut(float lo, float hi) 
{
  mPrimVtxVpdVzDiff[0] = lo;
  mPrimVtxVpdVzDiff[1] = hi;
}


//_________________
void StFlow::SetP(float lo, float hi) 
{
  mTrackP[0] = lo;
  mTrackP[1] = hi;
}


//_________________
void StFlow::SetPt(float lo, float hi) 
{
  mTrackPt[0] = lo;
  mTrackPt[1] = hi;
}


//_________________
void StFlow::SetTrackNHits(int lo, int hi) 
{
  mTrackNHits[0] = lo;
  mTrackNHits[1] = hi;  
}


//_________________
void StFlow::SetTrackNHitsFit(int lo, int hi) 
{
  mTrackNHitsFit[0] = lo;
  mTrackNHitsFit[1] = hi;
}


//_________________
void StFlow::SetTrackEta(float lo, float hi) 
{
  mTrackEta[0] = lo;
  mTrackEta[1] = hi;
}


//_________________
void StFlow::SetTrackFlag(short lo, short hi)
{
  mTrackFlag[0] = lo;
  mTrackFlag[1] = hi;
}


//_________________
void StFlow::SetTrackDCA(float lo, float hi)
{
  mTrackDcaGlobal[0] = lo;
  mTrackDcaGlobal[1] = hi;
}


//_________________
Bool_t StFlow::IsTofTrack(StMuTrack *trk)
{ //Only for globals
  Bool_t isTofHit = false;
  if(trk->btofPidTraits().beta() > 0 &&
     trk->btofPidTraits().timeOfFlight() > 0) 
  {
    isTofHit = true;
  }
  return isTofHit;
}
