#include "StHijingQA.h"
#include "StBTofHeader.h"
#include "TBranch.h"

ClassImp(StHijingQA)

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



StHijingQA::StHijingQA(const char *dir,
                       const char *fileName,
                       const char *filter,
                       int maxFiles,
                       const char *outFile)
{
  std::cout << "StHijingQA::StHijingQA - Creating an instance...";

  mEventsPassed = 0;

  mIsGoodTrack = false;
  mSGEvent = 0;
  mFileName = outFile;

  //
  // Initialize histograms variables
  //
  mRefMultMin  = 0;
  mRefMultMax  = 1000;
  mRefMultBins = 1000;

  mNTracksMin  = 0;
  mNTracksMax  = 1000;
  mNTracksBins = 1000;

  mPBins = 200;
  mPMin  = 0.;
  mPMax  = 2.1;

  mPtBins = 200;
  mPtMin  = 0.;
  mPtMax  = 2.1;

  mPxBins = 200;
  mPxMin  = 0.;
  mPxMax  = 2.1;

  mPyBins = 200;
  mPyMin  = 0.;
  mPyMax  = 2.1;

  mPzBins = 200;
  mPzMin  = 0.;
  mPzMax  = 2.1;

  mEtaBins = 300;
  mEtaMin  = -3;
  mEtaMax  = 3;

  //
  // Initialize single-particle cut variables
  //
  mTrackP[0] = 0.1;
  mTrackP[1] = 2.;
  mTrackEta[0] = -1.1;
  mTrackEta[1] = 1.1;

  InitChain(dir, fileName, filter, maxFiles);

  mNEvents = (unsigned int)mChain->GetEntries();
  std::cout << "*** NEvents to read = " << mNEvents << "***\n";

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



StHijingQA::~StHijingQA()
{
  /* nothing to do */
}



int StHijingQA::Init()
{
  std::cout << "StHijingQA::Init - Initializing the maker"
      << "--- EVENT OPTIONS ---" << std::endl
      << "\tmRefMultBins = " << mRefMultBins 
      << " mRefMultMin = " << mRefMultMin
      << " mRefMultMax = " << mRefMultMax << std::endl
      << "\tmNTracksBins = " << mNTracksBins
      << " mNTracksMin = " << mNTracksMin 
      << " mNTracksMax = " << mNTracksMax << std::endl;

  std::cout << "--- TRACK OPTIONS ---" << std::endl
      << "\tmPBins = " << mPBins
      << " mPMin = " << mPMin
      << " mPMax = " << mPMax << std::endl
      << "\tmPtBins = " << mPtBins
      << " mPtMin = " << mPtMin
      << " mPtMax = " << mPtMax
      << "\tmPxBins = " << mPxBins
      << " mPxMin = " << mPxMin
      << " mPxMax = " << mPxMax << std::endl
      << "\tmPyBins = " << mPyBins
      << " mPyMin = " << mPyMin
      << " mPyMax = " << mPyMax << std::endl
      << "\tmPzBins = " << mPzBins
      << " mPzMin = " << mPzMin
      << " mPzMax = " << mPzMax << std::endl
      << "\tmEtaBins = " << mEtaBins
      << " mEtaMin = " << mEtaMin
      << " mEtaMax = " << mEtaMax
      << std::endl;

  //
  // Create histograms and output file
  //
  mOutFile = new TFile(mFileName, "RECREATE");
  //
  // General event distributions
  //
  hRefMult = new TH1F("hRefMult", "Reference Multiplicity;RefMult;#",
                      mRefMultBins, mRefMultMin, mRefMultMax);
  hNTracks = new TH1F("hNTracks", ";N_{tracks} per event;#",
                      mNTracksBins, mNTracksMin, mNTracksMax);
  //
  // General track distributions
  //
  hP = new TH1F("hP", ";p (GeV/c);#", mPBins, mPMin, mPMax);
  hPt = new TH1F("hPt", ";p_{T} (GeV/c);#", mPtBins, mPtMin, mPtMax);
  hPx = new TH1F("hPx", ";p_{x} (GeV/c);#", mPxBins, mPxMin, mPxMax);
  hPy = new TH1F("hPy", ";p_{y} (GeV/c);#", mPyBins, mPyMin, mPyMax);
  hPz = new TH1F("hPz", ";p_{z} (GeV/c);#", mPzBins, mPzMin, mPzMax);
  hEta = new TH1F("hEta", ";#eta;#", mEtaBins, mEtaMin, mEtaMax);
  hEtaNoCut = new TH1F("hEtaNoCut", ";#eta;#", mEtaBins, mEtaMin, mEtaMax);
  hEtaVsPt = new TH2F("hEtaVsPt", ";#eta;p_{T} (GeV/c)",
                      250, -5., 5.,
                      400, 0., 2.);
  hVxVy = new TH2F("hVxVy", ";V_{x} (cm); V_{y} (cm)",
                   100, -1., 1.,
                   100, -1., 1.);
  hVyVz = new TH2F("hVyVz", ";V_{y} (cm); V_{z} (cm)",
                   100, -1., 1.,
                   100, -1., 1.);
  hVxVz = new TH2F("hVxVz", ";V_{x} (cm); V_{z} (cm)",
                   100, -1., 1.,
                   100, -1., 1.);
  hMassSqrVsPt = new TH2F("hMassSqrVsPt", ";p_{T}*Q (GeV/c);m^{2} (GeV/c^{2})",
                          400, -2., 2.,
                          1000, -0.1, 1.5);

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



int StHijingQA::InitChain(string dir, string fileName,
                          string filter, int maxFiles)
{
  mChain = new TChain("genevents", "genevents");

  std::cout << "StHijingQA[INFO]: TChain construct..." << std::endl;

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



int StHijingQA::FillChain(TChain *chain, char *dir,
                          const char *filter, int maxFiles)
{
  TSystem *gSystem = new TSystem();

  if (!gSystem)
  {
    std::cout << "StHijingQA[ERROR]: cannot allocate memory for TSystem\n";
    return 0;
  }

  void *pDir = gSystem->OpenDirectory(dir);
  if (!pDir)
  {
    std::cout << "StHijingQA[ERROR]: can't open directory " 
        << dir << std::endl;
    delete gSystem;
    return 0;
  }

  int count = 0;
  char *fileName;

  while ((fileName = (char *)gSystem->GetDirEntry(dir)))
  {
    if (strcmp(fileName, ".") == 0 || strcmp(fileName, "..") == 0)
      continue;

    if (strstr(fileName, filter))
    {
      char *fullName = gSystem->ConcatFileName(dir, fileName);
      std::cout << "StHijingQA[INFO]: Adding "
          << fullName << " to the chain" << std::endl;
      chain->Add(fullName);
      delete fullName;
      count++;
      if ((maxFiles > 0) && (count > maxFiles))
        break;
    }
  }

  std::cout << "StHijingQA[INFO]: Added " << count 
      << " files to the chain" << std::endl;

  delete gSystem;
  return count;
}


int StHijingQA::FillChain(TChain *chain, const char *fileName,
                          int maxFiles)
{
  std::ifstream inStream(fileName);

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

  std::cout << "StHijing[INFO]: using file list: "
      << fileName << std::endl;

  int count = 0;
  char buf[0xFF];

  for ( ; inStream.good(); )
  {
    inStream.getline(buf, 0xFF);
    TString sFileName(buf);
    if (sFileName.Contains("root"))
    {
      chain->Add(buf);
      count++;
      if ((maxFiles > 0) && (count > maxFiles))
        break;
    }
  }

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



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



int StHijingQA::Make()
{
  //   mSGEvent->Clear();
  int bytes = mChain->GetEntry(mEventsPassed++);

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

  if (!bytes)
  {
    std::cout << "StHijing[INFO]: no events left\n";
    return 0;
  }

  if (mSGEvent)
  {
    int nTracks = mSGEvent->GetNumberOfParticles();
    int refMult = 0;
    hNTracks->Fill(nTracks);
    for (int iTrack = 0; iTrack < nTracks; iTrack++)
    {
      StarGenParticle *particle = (*mSGEvent)[iTrack];

      //
      // Get particle data
      //
      const TLorentzVector &p = particle->momentum();
      float ptot = p.P();
      float px   = p.Px();
      float py   = p.Py();
      float pz   = p.Pz();
      if (ptot < 0.05 || 
          fabs(px) < 1e-5 || 
          fabs(py) < 1e-5 || 
          fabs(pz) < 1e-5)  
        continue;

      int   pdg    = particle->GetId();
      char  charge = pdg >= 0 ? 1 : -1;
      float mSqr = particle->GetMass();
      float vx   = particle->GetVx();
      float vy   = particle->GetVy();
      float vz   = particle->GetVz();
      float pt   = p.Pt();
      float eta  = p.Eta();

      //
      // Fill histograms
      //
      hP->Fill(ptot);
      hPt->Fill(pt);
      hPx->Fill(p.Px());
      hPy->Fill(p.Py());
      hPz->Fill(p.Pz());
      hEtaVsPt->Fill(eta, pt);
      hEtaNoCut->Fill(eta);
      hVxVy->Fill(vx, vy);
      hVxVz->Fill(vx, vz);
      hVyVz->Fill(vy, vz);

      if (ptot > mTrackP[0])
      {
        hMassSqrVsPt->Fill(pt*charge, mSqr);
        hEta->Fill(eta);
        if (eta > mTrackEta[0] && eta < mTrackEta[1])
          refMult++;
      }
    }
    hRefMult->Fill(refMult);
  }

  return kStOk;
}



int StHijingQA::Finish()
{

  std::cout << "*************************************" << std::endl
      << "StHijingQA has been finished" << std::endl
      << "\t nEventsProcessed: " << mNEvents << std::endl
      << "*************************************" << std::endl;
  mOutFile->cd();
  mOutFile->Write();
  mOutFile->Close();

  return kStOk;
}



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



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



unsigned int StHijingQA::GetNEvents()
{
  return mNEvents;
}
