#include "StQGSMdstQAMaker.h"

ClassImp(StQGSMdstQAMaker)

//_________________
StQGSMdstQAMaker::StQGSMdstQAMaker(const char *dirName,
                                   const char *fileName,
                                   const char *filter,
                                   int        maxFiles) {
  mDir = string(dirName);
  mFileName = string(fileName);
  mFilter = string(filter);

  mTree     = 0;
  mChain    = 0;
  mO97Event = 0;
  mEventId  = 0;
}


//_________________
Int_t StQGSMdstQAMaker::Init() {
  std::cout << "[StQGSMdstQAMaker]: Initialization... ";
  InitRead(mDir, mFileName, mFilter, mMaxFiles);
  mNEvents = (unsigned int)mChain->GetEntries();
  mOutFile = new TFile(mOutFileName, "RECREATE");
  if (!mOutFile) {
    std::cout << "FAIL" << std::endl;
    return kStErr;
  }
  else
    std::cout << "OK" << std::endl;
  //
  // Load PDG database
  //
  pdgDb = new TDatabasePDG;
  pdgDb->ReadPDGTable("StRoot/StFemtoDstMaker/pdg_table.txt");
  //
  // Creating event histograms
  //
  int refMultBins = 500;
  float refMultLo = 0.;
  float refMultHi = 500;
  hRefMultP90  = new TH1F("hRefMultP90", "reference multiplicity (greg) p_{#perp} > 0.09 GeV/c;RefMult;#",  refMultBins, refMultLo, refMultHi);
  hRefMultP100 = new TH1F("hRefMultP100", "reference multiplicity (greg) p_{#perp} > 0.1 GeV/c;RefMult;#",  refMultBins, refMultLo, refMultHi);
  hRefMultP120 = new TH1F("hRefMultP120", "reference multiplicity (greg) p_{#perp} > 0.12 GeV/c;RefMult;#", refMultBins, refMultLo, refMultHi);
  hRefMultP150 = new TH1F("hRefMultP150", "reference multiplicity (greg) p_{#perp} > 0.15 GeV/c;RefMult;#", refMultBins, refMultLo, refMultHi);
  hImpPar      = new TH1F("hImpPar", "impact parameter;imp (fm/c);#", 100, 0., 20.);
  hSph         = new TH1F("hSph",
                          "Transverse sphericity (|#eta| < 0.5, p_{#perp} > 0.15 GeV/c);S_{#perp},#frac{dN}{dS_{#perp}}",
                          100, 0., 1.);
  hSph2        = new TH1F("hSph2",
                          "Transverse sphericity (|#eta| < 1.0, p_{#perp} > 0.15 GeV/c);S_{#perp},#frac{dN}{dS_{#perp}}",
                          100, 0., 1.);
  hImpVsRef    = new TH2F("hImpVsRef", "p_{#perp} > 0.15 [GeV/c] && |#eta| < 0.5;b [fm];RefMult", 100, 0., 20., refMultBins, refMultLo, refMultHi);

  //
  // Creating track histograms
  //
  hVx = new TH1F("hVx", "x freeze-out position;V_{x} (fm);#", 100, -2., 2.);
  hVy = new TH1F("hVy", "y freeze-out position;V_{y} (fm);#", 100, -2., 2.);
  hVz = new TH1F("hVz", "z freeze-out position;V_{z} (fm);#", 100, -2., 2.);
  hPx = new TH1F("hPx", "single particle p_{x};p_{x} (GeV/c);#", 200, -2., 2.);
  hPy = new TH1F("hPy", "single particle p_{y};p_{y} (GeV/c);#", 200, -2., 2.);
  hPz = new TH1F("hPz", "single particle p_{z};p_{z} (GeV/c);#", 200, -2., 2.);
  hP  = new TH1F("hP", "single particle p;p (GeV/c);#", 400, 0., 2.);
  hPt = new TH1F("hPt", "single particle p_{T};p_{T} (GeV/c);#", 400, 0., 2.);  
  hEta = new TH1F("hEta", ";#eta;#", 200, -10., 10.);
  hEtaInt = new TH1F("hEtaInt", ";#eta;#", 200, -10., 10.);  
  hMSqrVsPt = new TH2F("hMSqrVsPt",
                       "m^{2} vs p_{T};p_{T}*charge (GeV/c);m^{2} (GeV^{2})",
                       200, -2., 2., // Pt
                       100, 0., 2.); // msqr
  hMomElectron = new TH1F("hMomElectron", "Momentum of e^{#pm};p (GeV/c);#frac{dN}{dp}",
                          250, 0., 2.5);
  hMomPion     = new TH1F("hMomPion", "Momentum of #pi^{#pm};p (GeV/c);#frac{dN}{dp}",
                          250, 0., 2.5);
  hMomKaon     = new TH1F("hMomKaon", "Momentum of K^{#pm};p (GeV/c);#frac{dN}{dp}",
                          250, 0., 2.5);
  hMomProton   = new TH1F("hMomProton", "Momentum of p^{#pm};p (GeV/c);#frac{dN}{dp}",
                          250, 0., 2.5);

  hRfrVsTfr_pion = new TH2F("hRfrVsTfr_pion", ";R^{#pi^{ch}}_{freeze out} [fm];#tau^{#pi^{ch}}_{freeze out} [fm/c]",
                            300, 0., 105.,
                            300, 0., 105.);
  hRfrVsTfr_kaon = new TH2F("hRfrVsTfr_kaon", ";R^{K^{ch}}_{freeze out} [fm];#tau^{K^{ch}}_{freeze out} [fm/c]",
                            300, 0., 105.,
                            300, 0., 105.);
  hRfrVsTfr_proton = new TH2F("hRfrVsTfr_proton", ";R^{p^{#pm}}_{freeze out} [fm];#tau^{p^{#pm}}_{freeze out} [fm/c]",
                              300, 0., 105.,
                              300, 0., 105.);

  hdNdTfr_pion = new TH1F("hdNdTfr_pion", ";#frac{dN}{d#tau^{#pi^{ch}_{freezeout}}};#tau^{#pi^{ch}}_{freezeout} [fm/c]",
                          105, 0., 105.);
  hdNdTfr_kaon = new TH1F("hdNdTfr_kaon", ";#frac{dN}{d#tau^{K^{ch}_{freezeout}}};#tau^{K^{ch}}_{freezeout} [fm/c]",
                          105, 0., 105.);
  hdNdTfr_proton = new TH1F("hdNdTfr_proton", ";#frac{dN}{d#tau^{p^{ch}_{freezeout}}};#tau^{p^{ch}}_{freezeout} [fm/c]",
                            105, 0., 105.);

  return kStOk;
}


//_________________
Int_t StQGSMdstQAMaker::Finish() {
  UninitRead();

  mOutFile->Write();
  mOutFile->Close();

  delete pdgDb;

  return kStOk;
}


//_________________
int StQGSMdstQAMaker::FillChain(TChain *chain, char *dir,
                                const char *filter, int maxFiles) {
  TSystem *gSystem = new TSystem();
  if (!gSystem) {
    std::cout << "[StQGSMdstQAMaker]: can't allocate memory for TSystem" << std::endl;;
    return 0;
  }

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

  int count = 0;
  char *file;

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

  std::cout << "[StQGSMdstQAMaker]: Added " << count 
            << " files to the chain" << std::endl;
  delete gSystem;
  return count;
}


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


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

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

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



//_________________
int StQGSMdstQAMaker::InitRead(string dir, string fileName,
                               string filter, int maxFiles) {
  mEventId = 0;

  mChain = new TChain("StO97Dst", "StO97Dst");

  std::cout << "[StQGSMdstQAMaker]: Call InitRead" << 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("StO97Event", &mO97Event);
  mNEvents = (unsigned int)mChain->GetEntries();
  return numFiles;
}


//_________________
void StQGSMdstQAMaker::UninitRead() {
  if (mChain)     delete mChain;
  mO97Event = 0;
  mChain    = 0;
}


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


//_________________
Int_t StQGSMdstQAMaker::Make() {
  if (!mNEvents) {
    std::cout << "[StQGSMdstQAMaker]: no events to read" << std::endl;
    return 0;
  }

  mO97Event->Clear();
  int bytes = mChain->GetEntry(mEventId++);

  if (mNEvents < mEventId) {
    std::cout << "[StQGSMdstQAMaker]: EOF" << std::endl;
    return 0;
  }

  if (!bytes) {
    std::cout << "[StQGSMdstQAMaker]: no event" << std::endl;
    return 0;
  }

  if (mO97Event) {
    int refMultP90 = 0;
    int refMultP100 = 0;
    int refMultP120 = 0;
    int refMultP150 = 0;
    float imp = mO97Event->GetImpactPar();

    TClonesArray *tracks = mO97Event->GetTracks();

    if (tracks) {
      int nTracks = tracks->GetEntries();
      for (int i = 0; i < nTracks; i++) {
        StO97Track *track = (StO97Track *)tracks->At(i);
        float px = track->GetPx();
        float py = track->GetPy();
        float pz = track->GetPz();
        float pt = sqrt(px*px + py*py);
        float p = sqrt(px*px + py*py + pz*pz);
        float massSqr = track->GetMass()*track->GetMass();
        float vx = track->GetXfr();
        float vy = track->GetYfr();
        float vz = track->GetZfr();
        float Rfr = TMath::Sqrt(vx*vx + vy*vy);
        float Tfr = track->GetTfr();
        float eta = pt > 0.05 /* i.e. pt > 0 */ ? 0.5*log((p + pz)/(p - pz)) : -999.;
        int   pdg = track->GetPdgId(pdgDb);
        int   charge = track->GetCharge(pdgDb, true);

        hEtaInt->Fill(eta);

        if (pt < 0.075 || fabs(eta) >= 0.5 /* STAR TPC acceptance */ ||
            charge == 0. /* only charged particles */ ||
            track->GetIsSpec()) {

          continue;
        }

        if (pt >= 0.09)   refMultP90++;
        if (pt >= 0.10)   refMultP100++;
        if (pt >= 0.12)   refMultP120++;
        if (pt >= 0.15)   refMultP150++;

        if (pt >= 0.1 /* STAR TPC lower momentum limit from STAR TPC NIM */) {
          hMSqrVsPt->Fill(pt*charge, massSqr);
          hVx->Fill(vx);
          hVy->Fill(vy);
          hVz->Fill(vz);
          hPx->Fill(px);
          hPy->Fill(py);
          hPz->Fill(pz);
          hPt->Fill(pt);
          hP->Fill(p);
          hEta->Fill(eta);
        }

        switch (abs(pdg)) {
        case 11: // Electron
          hMomElectron->Fill(p);
          break;
        case 211: // Pion
          hMomPion->Fill(p);
          hdNdTfr_pion->Fill(Tfr);
          hRfrVsTfr_pion->Fill(Rfr, Tfr);
          break;
        case 321: // Kaon
          hMomKaon->Fill(p);
          hdNdTfr_kaon->Fill(Tfr);
          hRfrVsTfr_kaon->Fill(Rfr, Tfr);
          break;
        case 2212: // Proton
          hMomProton->Fill(p);
          hdNdTfr_proton->Fill(Tfr);
          hRfrVsTfr_proton->Fill(Rfr, Tfr);
        } // switch (abs(pdg))
      } // for (int i = 0; i < nTracks; i++)
    } // if (tracks)
    else {
      std::cout << "[StQGSMdstQAMaker]: tracks == NULL" << std::endl;
    }

    if (refMultP90)
      hRefMultP90->Fill(refMultP90);
    if (refMultP100)
      hRefMultP100->Fill(refMultP100);
    if (refMultP120)
      hRefMultP120->Fill(refMultP120);
    if (refMultP150)
      hRefMultP150->Fill(refMultP150);

    hSph->Fill(mO97Event->GetTransverseSphericity());
    hSph2->Fill(mO97Event->GetTransverseSphericity2());

    hImpPar->Fill(imp);
    hImpVsRef->Fill(imp, refMultP150);
  } // if (mO97Event)

  return kStOk;
}