#include "StHbtStarGenReader.h"
#include "TString.h"

ClassImp(StHbtStarGenReader)

  //_________________
  StHbtStarGenReader::StHbtStarGenReader(const char *dir, 
                                         const char *fileName,
                                         const char *filter,
                                         int maxFiles,
                                         bool debug) {
    mDir      = string(dir);
    mFileName = string(fileName);
    mFilter   = filter;
    mMaxFiles = maxFiles;
    mDebug    = debug;

    mTotalTracks = 0.;

    mTree       = 0;
    mTChain     = 0;
    mSGEvent    = 0;
    mEventIndex = 0;

    randy = new TRandom3(0);

    InitRead(mDir, mFileName, mFilter, mMaxFiles);
    mNEvents = (unsigned int)mTChain->GetEntries();
  }

//_________________
int StHbtStarGenReader::FillChain(TChain *chain, char *dir,
                                  const char *filter, int maxFiles) {

  TSystem *gSystem = new TSystem();
  if (!gSystem) {
    std::cout << "StHbtStarGenReader[ERROR]: cannot allocate memory for TSystem\n";
    return 0;
  }
  void *lDir = gSystem->OpenDirectory(dir);
  if (!lDir) {
    std::cout << "StHbtStarGenReader[ERROR]: can't open directory " 
        << dir << std::endl;
    delete gSystem;
    return 0;
  }

  int mCount = 0;
  char *lFileName;

  while ((lFileName = (char *)gSystem->GetDirEntry(dir))) {
    if (strcmp(lFileName, ".") == 0 || strcmp(lFileName, "..") == 0)
      continue;
    if (strstr(lFileName, filter)) {
      char *lFullName = gSystem->ConcatFileName(dir, lFileName);
      std::cout << "StHbtStarGenReader[INFO]: Adding "
          << lFullName << " to the chain" << std::endl;
      chain->Add(lFullName);
      delete lFullName;
      mCount++;
      if ((maxFiles > 0) && (mCount > maxFiles)) break;
    }
  }

  std::cout << "StHbtStarGenReader[INFO]: Added " << mCount 
      << " files to the chain" << std::endl;
  delete gSystem;
  return mCount;
}

//_________________
int StHbtStarGenReader::FillChain(TChain *chain, const char *fileName,
                                  int maxFiles) {

  std::ifstream lInStream(fileName);

  if (!lInStream.is_open()) {
    std::cout << "StHbtStarGenReader[ERROR]: can't open file list: "
        << fileName << std::endl;
    return 0;
  }
  std::cout << "StHbtStarGenReader[INFO]: using file list: "
      << fileName << std::endl;

  int mCount = 0;
  char buf[0xFF];
  for ( ; lInStream.good(); ) {
    lInStream.getline(buf, 0xFF);
    TString lFileName(buf);
    if (lFileName.Contains("root")) {
      chain->Add(buf);
      mCount++;
      if ((maxFiles > 0) && (mCount > maxFiles)) break;
    }
  }

  std::cout << "StHbtStarGenReader[INFO]: Added " << mCount
      << " files to the chain" << std::endl;
  return mCount;
}

//_________________
int StHbtStarGenReader::InitRead(string aDir, string fileName,
                                 string aFilter, int maxFiles) {
  mEventIndex = 0;
  //  mSGEvent = new StSGEvent();
  mTChain = new TChain("genevents", "genevents");

  std::cout << "StHbtStarGenReader[INFO]: Call InitRead" << std::endl;

  int lNumFiles = 0;
  if (!fileName.empty()) {
    if (strstr(fileName.c_str(), ".lis") ||
        strstr(fileName.c_str(), ".list")) {
      lNumFiles = FillChain( mTChain, 
                            (aDir + fileName).c_str(), 
                            maxFiles );
    }
    else {
      mTChain->Add((aDir + fileName).c_str());
      lNumFiles++;
    }
  } else {
    lNumFiles = FillChain(mTChain, (char *)aDir.c_str(), // not cool
                          aFilter.c_str(), maxFiles);
  }
  mTChain->SetBranchAddress("primaryEvent", &mSGEvent);
  mNEvents = (unsigned int)mTChain->GetEntries();
  std::cout << "StHbtStarGenReader::InitRead [INFO]"
      << " NEvents to read: " << mNEvents << std::endl;
  return lNumFiles;
}

//_________________
void StHbtStarGenReader::UninitRead()
{
  //  if (mSGEvent) delete mSGEvent;
  if (mTChain) delete mTChain;
  mSGEvent = 0;
  mTChain   = 0;
}

//_________________
void StHbtStarGenReader::PrintTotalTracks()
{
  std::cout << "StHbtStarGenReader[INFO]: mTotalTracks = "
      << mTotalTracks << std::endl;
}

//_________________
int StHbtStarGenReader::GetNEvents()
{
  if (mTChain) {
    return mNEvents;
  }
  else {
    return -1;
  }
}

//_________________
StHbtEvent *StHbtStarGenReader::Read()
{
  if(mDebug)
  {
    std::cout << "StHbtStarGenReader[DEBUG]: mEventIndex = "
        << mEventIndex << " mNEvents = "
        << mNEvents << std::endl;
  }

  if(!mNEvents) {
    std::cout << "StHbtStarGenReader[INFO]: no events to read\n";
    return 0;
  }

  int lBytes = mTChain->GetEntry(mEventIndex++);

  if(mNEvents < mEventIndex) {
    std::cout << "StHbtStarGenReader[INFO]: EOF\n";
    return 0;
  }

  if(!lBytes) {
    std::cout << "StHbtStarGenReader[INFO]: no event\n";
    return 0;
  }

  StHbtEvent *mHbtEvent = NULL;

  if (mSGEvent)
  {
    mHbtEvent = new StHbtEvent(); // create new StHbtEvent
    const float magRFF = -4.98;   // [10^3 G], ref: p+Au 200 GeV y2015
    float Vz = 0.0; // randy->Uniform(-30., 30.);
    float Vx = 0.0; // randy->Uniform(-0.5, 0.5);
    float Vy = 0.0; // randy->Uniform(-0.5, 0.5);
    StThreeVectorF v(Vx, Vy, Vz); // event vertex
    int nTracks = mSGEvent->GetNumberOfParticles(); // number of tracks in an event

    //
    // Set an event information.
    //
    //  Number of the positive and negative tracks
    //  are setting after the loop over tracks.
    //
    mHbtEvent->SetEventNumber(mSGEvent->GetEventNumber());
    mHbtEvent->SetRunNumber(mSGEvent->GetRunNumber());
    mHbtEvent->SetZdcAdcEast(0.);
    mHbtEvent->SetZdcAdcWest(0.);
    mHbtEvent->SetUncorrectedNumberOfPrimaries(nTracks);
    mHbtEvent->SetPrimVertPos(v);
    mHbtEvent->SetMagneticField(magRFF);
    mHbtEvent->SetVpdVz(0.);
    mHbtEvent->SetCent9(0);
    mHbtEvent->SetCent16(0);
    mHbtEvent->SetRefMultCorr(nTracks);
    mHbtEvent->SetSimulation(true);

    mTotalTracks += (float)nTracks;
    if (mDebug)
      std::cout << "StHbtStarGenReader[DEBUG]: nTracks = "
          << nTracks << std::endl;

    int nTracksPos = 0;
    int nTracksNeg = 0;

    for (int i = 0; i < nTracks; i++)
    {
      StarGenParticle *sParticle = (*mSGEvent)[i];


      //
      // Calculate charge of the particle
      //
      const TLorentzVector &p = sParticle->momentum();
      float ptot = p.P();
      float px   = p.Px();
      float py   = p.Py();
      float pz   = p.Pz();
      if (fabs(px) < 1e-5 ||
          fabs(py) < 1e-5 ||
          fabs(pz) < 1e-5)
        continue;
      int pdg = sParticle->GetId();
      int charge;
      float eta = p.Eta();
      float pt  = p.Pt();

      if (eta <= -1.1 || eta >= 1.1 || pt <= 0.1)
        continue;

      StHbtTrack *mHbtTrack = new StHbtTrack;
      if (pdg >= 0)
      {
        nTracksPos++;
        charge = +1;
      }
      else
      {
        nTracksNeg++;
        charge = -1;
      }

      mHbtTrack->SetTrackId(i);
      mHbtTrack->SetFlag(300);      // old flag < 1000 means no pile up

      //
      // Energy loss of the particle
      //
      int nHits = 45;

      mHbtTrack->SetNHits(nHits*charge);
      mHbtTrack->SetNHitsPossible(nHits);
      mHbtTrack->SetNHitsDedx(nHits);
      mHbtTrack->SetdEdx(dedxMean(sParticle->GetMass(), p.P()));
      mHbtTrack->SetChi2(1.);

      //
      // Calculate nSigma of the
      //  electron, pion, kaon, proton
      //
      float nSigmaElectron = -999.;
      float nSigmaPion = -999.;
      float nSigmaKaon = -999.;
      float nSigmaProton = -999.;

      switch (abs(pdg))
      {
        case 211:
          nSigmaPion = 0.;
          break;
        case 321:
          nSigmaKaon = 0.;
          break;
        case 2212:
          nSigmaProton = 0.;
          break;
      }

      mHbtTrack->SetNSigmaElectron(nSigmaElectron);
      mHbtTrack->SetNSigmaPion(nSigmaPion);
      mHbtTrack->SetNSigmaKaon(nSigmaKaon);
      mHbtTrack->SetNSigmaProton(nSigmaProton);

      //
      // Calculate distance of closet approach
      //
      float xDCA = sParticle->GetVx();
      float yDCA = sParticle->GetVy();
      float zDCA = sParticle->GetVz();
      float xyDCA = sqrt(xDCA*xDCA + yDCA*yDCA);

      mHbtTrack->SetDCAxyGlobal(xyDCA);
      mHbtTrack->SetDCAzGlobal(zDCA);

      //
      // Set momentum of the particle
      //
      StThreeVectorF mGlobMomentum(p.Px(),
                                   p.Py(),
                                   p.Pz());
      StThreeVectorF mPrimMomentum(mGlobMomentum);
      mHbtTrack->SetP(mPrimMomentum);

      if(mDebug)
        std::cout << Form("Track parameters: px=%3.2f py=%3.2f pz=%3.2f q=%1d nhits=%2d",
                          mPrimMomentum.x(), mPrimMomentum.y(), mPrimMomentum.z(),
                          charge, nHits) << std::endl;

      //
      // Calculate helix
      //
      StThreeVectorF mGlobOrigin(xDCA, yDCA, zDCA);
      StThreeVectorF mPrimOrigin(0., 0., 0.);
      StPhysicalHelixD mGlobHelix(mGlobMomentum, mGlobOrigin, magRFF*kilogauss, charge);
      StPhysicalHelixD mPrimHelix(mPrimMomentum, mPrimOrigin, magRFF*kilogauss, charge);

      mHbtTrack->SetHelixGlobal(mGlobHelix);
      mHbtTrack->SetHelix(mPrimHelix);

      //
      // Calculate topology map
      //
      unsigned int tmap1 = 0;
      unsigned int tmap2 = 0;
      if (ptot < 0.15) // tracks only in the inner sector. R=1m
      {
        tmap1 = 0;
        tmap2 = 0;
        for (int iBit = 0; iBit < 32; iBit++)
        {
          if (randy->Rndm() < 0.95)
            tmap1 |= 1 << iBit;
          else
            tmap1 |= 0 << iBit;
        }
        tmap1 |= 1 << 0; // default set to the first bit: primary-vertex-used
      }
      else
      {
        if (ptot < 0.3) // tracks may come to the end of the outer sector: R = 2m
        { 
          tmap1 = 0;
          tmap2 = 0;
          for (int iBit = 0; iBit < 32; iBit++)
          {
            if (randy->Rndm() < 0.95)
              tmap1 |= 1 << iBit;
            else 
              tmap1 |= 0 << iBit;

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

            if(iBit < 27)
            {
              if(randy->Rndm() < 0.95) 
                tmap2 |= 1 << iBit;
              else 
                tmap2 |= 0 << iBit;
            }
          }
          tmap1 |= 1 << 0; // default set to the first bit: primary-vertex-used
        }
      }
      mHbtTrack->SetTopologyMap(0, tmap1);
      mHbtTrack->SetTopologyMap(1, tmap2);

      mHbtTrack->SetTofBeta(ptot/p.Energy());
      mHbtTrack->SetHiddenInfo(0);

      mHbtEvent->TrackCollection()->push_back(mHbtTrack);
    } // end of a loop over tracks

    //
    // Number of charged tracks 
    //
    mHbtEvent->SetUncorrectedNumberOfPrimaries(nTracksPos + nTracksNeg);
    mHbtEvent->SetUncorrectedNumberOfPositivePrimaries(nTracksPos);
    mHbtEvent->SetUncorrectedNumberOfNegativePrimaries(nTracksNeg);

  } // if event != 0
  else
  {
    //
    // FemtoDstMaker returned NULL, so we are not OK
    //
    mReaderStatus = 1;
  }

  return mHbtEvent;
}

//_________________
StHbtEvent *StHbtStarGenReader::ReturnHbtEvent() {

  StHbtEvent *mHbtEvent = Read();
  if (!mHbtEvent) {
    std::cout << "StHbtStarGenReader[WARN]: no hbt event (mHbtEvent == NULL)\n";
    return 0;
  }
  else {
    return mHbtEvent;
  }
}

//_________________
StHbtString StHbtStarGenReader::Report() {

  StHbtString temp = "\nThis is the StHbtStarGenReader\n";
  return temp;
}

//_________________
void StHbtStarGenReader::DoDebug(bool aDebug)
{
  mDebug = aDebug;
}

//________________
float StHbtStarGenReader::dedxMean(float mass, float momentum)
{
  Double_t dedxMean;
  Double_t tpcDedxGain = 0.174325e-06;
  Double_t tpcDedxOffset = -2.71889; 
  Double_t tpcDedxRise = 776.626;

  Double_t gamma = TMath::Sqrt(momentum*momentum/(mass*mass)+1.);
  Double_t beta = TMath::Sqrt(1. - 1./(gamma*gamma));
  Double_t rise = tpcDedxRise* beta*beta*gamma*gamma;

  if (beta > 0)
    dedxMean = tpcDedxGain/(beta*beta)*(0.5*TMath::Log(rise) - beta*beta 
                                        - tpcDedxOffset);
  else
    dedxMean = 1000.;

  return dedxMean;
}
