#include "StFemtoDstQAMaker.h"
#include "StThreeVector.hh"
#include "StThreeVectorF.hh"
#include "StPhysicalHelixD.hh"

ClassImp(StFemtoDstQAMaker)

  //_________________
  StFemtoDstQAMaker::StFemtoDstQAMaker(const char *aDir, 
                                       const char *aFileName,
                                       const char *aFilter,
                                       int aMaxFiles) {
    dir      = string(aDir);
    fileName = string(aFileName);
    filter   = aFilter;
    maxFiles = aMaxFiles;
    outFileName = "oFemtoQA.root";

    tree       = 0;
    chain      = 0;
    femtoEvent = 0;
    eventId    = 0;
    mTrig      = 0;

    mRunIdBins = 7164020;    //Total amount of runs in between Run16 and Run9
    mRunIdVal[0] = 10016001; //Beginning of Run9
    mRunIdVal[1] = 17180021; //End of Run16
  }

//_________________
Int_t StFemtoDstQAMaker::Init() {

  InitRead(dir, fileName, filter, maxFiles);
  nEvents = (unsigned int)chain->GetEntries();
  outFile = new TFile(outFileName, "RECREATE");

  Int_t mRefMultWeightBins = 200;
  Double_t mRefMultWeightLo = 0.;
  Double_t mRefMultWeightHi = 20.;
  Int_t mRefMultBins = 700;
  Double_t mRefMultLo = -0.5;
  Double_t mRefMultHi = 699.5;
  Int_t mPrimVtxZBins = 400;
  Double_t mPrimVtxZLo = -200.;
  Double_t mPrimVtxZHi = 200.;
  Int_t mPrimVertRBins = 60;
  Double_t mPrimVertRLo = -0.5;
  Double_t mPrimVertRHi = 5.5;
  Int_t mPrimVertXBins = 100;
  Double_t mPrimVertXLo = -5.;
  Double_t mPrimVertXHi = 5.;
  Int_t mPrimVertYBins = 100;
  Double_t mPrimVertYLo = -5.;
  Double_t mPrimVertYHi = 5.;
  Int_t mNSigmaBins = 200;
  Double_t mNSigmaLo = -10.;
  Double_t mNSigmaHi = 10.;
  Int_t mMomBins = 200;
  Double_t mMomLo = 0.;
  Double_t mMomHi = 2.;
  Int_t mGlobCoordBins = 400;
  Double_t mGlobCoordLo = -200.;
  Double_t mGlobCoordHi = 200.;
  Int_t mDedxBins = 200;
  Double_t mDedxLo = 0.;
  Double_t mDedxHi = 20.;
  Int_t mDCAxyBins = 500;
  Double_t mDCAxyLo = 0.;
  Double_t mDCAxyHi = 5.;
  Int_t mDCAzBins = 200;
  Double_t mDCAzLo = -50.;
  Double_t mDCAzHi = 50.;
  Int_t mMSqrBins = 280;
  Double_t mMSqrLo = -0.2;
  Double_t mMSqrHi = 1.2;
  Int_t mInvBetaBins = 180;
  Double_t mInvBetaLo = 0.6;
  Double_t mInvBetaHi = 2.4;
  Int_t mBetaBins = 100;
  Double_t mBetaLo = 0.;
  Double_t mBetaHi = 1.;
  Int_t mBTofMultBins = 700;
  Double_t mBTofMultLo = -0.5;
  Double_t mBTofMultHi = 699.5;

  // Event histograms
  hRefMultWeight = new TH1F("hRefMultWeight",";RefMultCorrWeight;#",
                            mRefMultWeightBins, mRefMultWeightLo, mRefMultWeightHi);
  hRefMult = new TH1F("hRefMult",
                      ";RefMult;#",
                      mRefMultBins, mRefMultLo, mRefMultHi);
  hRefMult2 = new TH1F("hRefMult2",
                       "|#eta| < 1;RefMult2;#",
                       mRefMultBins, mRefMultLo, mRefMultHi);
  hRefMultCorr = new TH1F("hRefMultCorr", ";RefMultCorr;#",
                          mRefMultBins, mRefMultLo, mRefMultHi);
  hTofMult = new TH1F("hTofMult", ";TofMult;#",
                      mRefMultBins, mRefMultLo, mRefMultHi);
  hTofVsRef = new TH2F("hTofVsRef", ";TofMult;RefMult",
                       mRefMultBins, mRefMultLo, mRefMultHi,
                       mRefMultBins, mRefMultLo, mRefMultHi);
  hTofVsRef2 = new TH2F("hTofVsRef2", ";TofMult;RefMult2",
                        mRefMultBins, mRefMultLo, mRefMultHi,
                        mRefMultBins, mRefMultLo, mRefMultHi);
  hRefVsRef2 = new TH2F("hRefVsRef2", ";RefMult;RefMult2",
                        mRefMultBins, mRefMultLo, mRefMultHi,
                        mRefMultBins, mRefMultLo, mRefMultHi);
  hPrimVertZ = new TH1F("hPrimVertZ",
                        "Z coordinate of primary vertex;z (cm);#",
                        mPrimVtxZBins, mPrimVtxZLo, mPrimVtxZHi);
  hPrimVertR = new TH1F("hPrimVertR","#sqrt{V_{x}^{2}+V_{y}^{2}} of primary vertex;R (cm);#",
                        mPrimVertRBins, mPrimVertRLo, mPrimVertRHi);
  hPrimVertXvsY = new TH2F("hPrimVertXvsY","X vs Y for primary vertex;x (cm);y (cm)",
                           mPrimVertXBins, mPrimVertXLo, mPrimVertXHi,
                           mPrimVertYBins, mPrimVertYLo, mPrimVertYHi);
  hPrimVertVpdVzDiff = new TH1F("hPrimVertVpdVzDiff","z_{TPC} - z_{VPD};z_{TPC} - z_{VPD} (cm);#",
                                300, -15., 15.);
  hRefMultVsRunNumer = new TProfile("hRefMultVsRunNumer","refMult vs. Run Number;Run Number;<refMult>",
                                    mRunIdBins, mRunIdVal[0], mRunIdVal[1]);
  hPrimTrackNumVsRunNumber = new TProfile("hPrimTrackNumVsRunNumber","pTrk vs. Run Number;Run Number;<PrimTrackNum>",
                                          mRunIdBins, mRunIdVal[0], mRunIdVal[1]);
  hGlobTrackNumVsRunNumber = new TProfile("hGlobTrackNumVsRunNumber","gTrk vs. Run Number;Run Number;<GlobTrackNum>",
                                          mRunIdBins, mRunIdVal[0], mRunIdVal[1]);
  hTofMultVsRunNumber = new TProfile("hTofMultVsRunNumber","TofMult vs. Run Number;Run Number;<TofMult>",
                                     mRunIdBins, mRunIdVal[0], mRunIdVal[1]);
  // Track histograms
  hNSigmaPionVsMom = new TH2F("hNSigmaPionVsMom", "N#sigma(#pi) vs momenta;p (GeV/c);N#sigma(pi)",
                              mMomBins,mMomLo,mMomHi,
                              mNSigmaBins, mNSigmaLo, mNSigmaHi);
  hNSigmaKaonVsMom = new TH2F("hNSigmaKaonVsMom", "N#sigma(K) vs momenta;p (GeV/c);N#sigma(K)",
                              mMomBins,mMomLo,mMomHi,
                              mNSigmaBins, mNSigmaLo, mNSigmaHi);
  hNSigmaProtonVsMom = new TH2F("hNSigmaProtonVsMom", "N#sigma(p) vs momenta;p (GeV/c);N#sigma(p)",
                                mMomBins,mMomLo,mMomHi,
                                mNSigmaBins, mNSigmaLo, mNSigmaHi);
  hNSigmaElectronVsMom = new TH2F("hNSigmaElectronVsMom", "N#sigma(e) vs momenta;p (GeV/c);N#sigma(e)",
                                  mMomBins,mMomLo,mMomHi,
                                  mNSigmaBins, mNSigmaLo, mNSigmaHi);

  hP = new TH1F("hP","Momentum;p (GeV/c);#",mMomBins,mMomLo,mMomHi);
  hPt = new TH1F("hPt","Transverse momentum;p_{T} (GeV/c);#",mMomBins,mMomLo,mMomHi);

  hGlobPx = new TH1F("hGlobPx",
                     "global track p_{x};p_{x} (GeV/c);#",
                     2*mMomBins,-mMomHi,mMomHi);
  hGlobPy = new TH1F("hGlobPy",
                     "global track p_{y};p_{y} (GeV/c);#",
                     2*mMomBins,-mMomHi,mMomHi);
  hGlobPz = new TH1F("hGlobPz",
                     "global track p_{z};p_{z} (GeV/c);#",
                     2*mMomBins,-mMomHi,mMomHi);

  hGlobX = new TH1F("hGlobX",
                    "global track first point x-coordinate;x (cm);#",
                    mGlobCoordBins, mGlobCoordLo, mGlobCoordHi);
  hGlobY = new TH1F("hGlobY",
                    "global track first point y-coordinate;y (cm);#",
                    mGlobCoordBins, mGlobCoordLo, mGlobCoordHi);
  hGlobZ = new TH1F("hGlobZ",
                    "global track first point z-coordinate;z (cm);#",
                    mGlobCoordBins, mGlobCoordLo, mGlobCoordHi);


  hdEdx = new TH1F("hdEdx", "dE/dx;dE/dx (keV/cm);N", mDedxBins, mDedxLo, mDedxHi);
  hdEdx_p = new TH2F("hdEdx_p", "dE/dx vs momentum;q*p (GeV/c); dE/dx (keV/cm);", 
                     2*mMomBins, -mMomHi, mMomHi, 
                     mDedxBins, mDedxLo, mDedxHi);
  hDCAxy = new TH1F("hDCAxy",
                    "Distance of closest approach (XY plane) for primary track;DCA_{xy} (cm);#",
                    mDCAxyBins, mDCAxyLo, mDCAxyHi);
  hDCAz = new TH1F("hDCAz",
                   "Distance of closest approach (Z axis) for primary track;DCA_{z} (cm);#",
                   mDCAzBins, mDCAzLo, mDCAzHi);
  hDCAxyGlob = new TH1F("hDCAxyGlob",
                        "Distance of closest approach (XY plane) for global track;DCA_{xy} (cm);#",
                        mDCAxyBins, mDCAxyLo, mDCAxyHi);
  hDCAzGlob = new TH1F("hDCAzGlob",
                       "Distance of closest approach (Z axis) for global track;DCA_{z} (cm);#",
                       mDCAzBins, mDCAzLo, mDCAzHi);
  hMSqr_P = new TH2F("hMSqr_P","m^{2} vs p;p (GeV/c);m^{2} (GeV/c^{2})^{2}",
                     2*mMomBins, -mMomHi, mMomHi, 
                     mMSqrBins, mMSqrLo, mMSqrHi);
  hInvBeta_P = new TH2F("hInvBeta_P",
                        "1/#beta vs p*q;q*p (GeV/c);1/#beta",
                        2*mMomBins, -mMomHi, mMomHi, 
                        mInvBetaBins, mInvBetaLo, mInvBetaHi);
  hBeta_P = new TH2F("hBeta_P",
                     "#beta vs momentum;q*p (GeV/c);#beta",
                     2*mMomBins, -mMomHi, mMomHi, 
                     mBetaBins, mBetaLo, mBetaHi);
  hEta = new TH1F("hEta",
                  "#eta;#eta;N",
                  100, -1., 1.);

  return kStOk;
}

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

  outFile->Write();
  outFile->Close();

  return kStOk;
}

//_________________
int StFemtoDstQAMaker::FillChain(TChain *aChain, char *aDir,
                                 const char *aFilter, int aMaxFiles) {

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

  int lCount = 0;
  char *lFileName;

  while ((lFileName = (char *)gSystem->GetDirEntry(aDir))) {
    if (strcmp(lFileName, ".") == 0 || strcmp(lFileName, "..") == 0)
      continue;
    if (strstr(lFileName, aFilter)) {
      char *lFullName = gSystem->ConcatFileName(aDir, lFileName);
      std::cout << "StFemtoDstQAMaker[INFO]: Adding "
          << lFullName << " to the chain" << std::endl;
      aChain->Add(lFullName);
      delete lFullName;
      lCount++;
      if ((aMaxFiles > 0) && (lCount > aMaxFiles)) break;
    }
  }

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

//_________________
int StFemtoDstQAMaker::FillChain(TChain *aChain, const char *aFileName,
                                 int aMaxFiles) {
  std::ifstream lInStream(aFileName);


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

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

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

//_________________
int StFemtoDstQAMaker::InitRead(string aDir, string aFileName,
                                string aFilter, int aMaxFiles) {
  eventId = 0;
  //  femtoEvent = new StFemtoEvent();
  chain = new TChain("StFemtoDst", "StFemtoDst");

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

  int lNumFiles = 0;
  if (!aFileName.empty()) {
    if (strstr(aFileName.c_str(), ".lis") ||
        strstr(aFileName.c_str(), ".list")) {
      lNumFiles = FillChain(chain, 
                            (aDir + aFileName).c_str(), 
                            aMaxFiles);
    }
    else {
      chain->Add((aDir + aFileName).c_str());
      lNumFiles++;
    }
  } else {
    lNumFiles = FillChain(chain, (char *)aDir.c_str(), // not cool
                          aFilter.c_str(), aMaxFiles);
  }
  chain->SetBranchAddress("StFemtoEvent", &femtoEvent);
  nEvents = (unsigned int)chain->GetEntries();
  return lNumFiles;
}

//_________________
void StFemtoDstQAMaker::UninitRead() {
  if (chain)     delete chain;
  femtoEvent = 0;
  chain   = 0;
}

//_________________
void StFemtoDstQAMaker::AddTrigId(unsigned int aTrigId) {
  mTrigIdCol.push_back(aTrigId);
}

//_________________
int StFemtoDstQAMaker::GetNEvents() {
  if (chain)
    return nEvents;
  else
    return -1;
}

//_________________
Int_t StFemtoDstQAMaker::Make() {
  if (!nEvents) {
    std::cout << "StFemtoDstQAMaker[INFO]: no events to read\n";
    return 0;
  }

  femtoEvent->Clear();
  int lBytes = chain->GetEntry(eventId++);

  if (nEvents < eventId) {
    std::cout << "StFemtoDstQAMaker[INFO]: EOF\n";
    return 0;
  }

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

  if (femtoEvent) {
    if (!mTrigIdCol.empty()) {
      unsigned int lTrigId = femtoEvent->GetTriggerId();
      if (((lTrigId & mTrig) == 0) || ((lTrigId & (!mTrig)) != 0))
        return 0;
    }

    // fill histograms
    hRefMultWeight->Fill(femtoEvent->GetRefMultCorrWeight());
    hRefMult->Fill(femtoEvent->GetRefMult());
    hRefMult2->Fill(femtoEvent->GetRefMult2());
    hRefMultCorr->Fill(femtoEvent->GetRefMultCorr());
    hTofMult->Fill(femtoEvent->GetNumberOfBTofHit());
    hTofVsRef->Fill(femtoEvent->GetNumberOfBTofHit(),
                    femtoEvent->GetRefMult());
    hTofVsRef2->Fill(femtoEvent->GetNumberOfBTofHit(),
                     femtoEvent->GetRefMult2());
    hRefVsRef2->Fill(femtoEvent->GetRefMult(),
                     femtoEvent->GetRefMult2());
    hPrimVertZ->Fill(femtoEvent->GetVertexPositionZ());
    hPrimVertR->Fill(femtoEvent->GetVertexPositionR());
    hPrimVertXvsY->Fill(femtoEvent->GetVertexPositionX(),
                        femtoEvent->GetVertexPositionY());
    hPrimVertVpdVzDiff->Fill(femtoEvent->GetVertexPositionZ() - femtoEvent->GetVpdVz());
    hRefMultVsRunNumer->Fill(femtoEvent->GetRunId(),femtoEvent->GetRefMult());
    hPrimTrackNumVsRunNumber->Fill(femtoEvent->GetRunId(),
                                   femtoEvent->GetNumberOfPrimaryTracks());
    hGlobTrackNumVsRunNumber->Fill(femtoEvent->GetRunId(),
                                   femtoEvent->GetNumberOfGlobalTracks());
    hTofMultVsRunNumber->Fill(femtoEvent->GetRunId(),
                              femtoEvent->GetNumberOfBTofHit());

    TClonesArray *lTracks = femtoEvent->GetFemtoTracks();

    if (lTracks) {
      int lNTracks = lTracks->GetEntries();
      for (int i = 0; i < lNTracks; i++) {
        StFemtoTrack *lTrack = (StFemtoTrack *)lTracks->At(i);

        hNSigmaPionVsMom->Fill(lTrack->GetCharge()*lTrack->GetPtot(), lTrack->GetNSigmaPion());
        hNSigmaKaonVsMom->Fill(lTrack->GetCharge()*lTrack->GetPtot(), lTrack->GetNSigmaKaon());
        hNSigmaProtonVsMom->Fill(lTrack->GetCharge()*lTrack->GetPtot(), lTrack->GetNSigmaProton());
        hNSigmaElectronVsMom->Fill(lTrack->GetCharge()*lTrack->GetPtot(), lTrack->GetNSigmaElectron());

        hGlobPx->Fill( lTrack->GetGlobalPx() );
        hGlobPy->Fill( lTrack->GetGlobalPy() );
        hGlobPz->Fill( lTrack->GetGlobalPz() );
        hGlobX->Fill( femtoEvent->GetVertexPositionX() + lTrack->GetDCAx() );
        hGlobY->Fill( femtoEvent->GetVertexPositionY() + lTrack->GetDCAy() );
        hGlobZ->Fill( femtoEvent->GetVertexPositionZ() + lTrack->GetDCAz() );

        StThreeVectorF vMomGlob( lTrack->GetGlobalPx(),
                                lTrack->GetGlobalPy(),
                                lTrack->GetGlobalPz() );
        StThreeVectorF vOrigGlob( femtoEvent->GetVertexPositionX() + lTrack->GetDCAx(),
                                 femtoEvent->GetVertexPositionY() + lTrack->GetDCAy(),
                                 femtoEvent->GetVertexPositionZ() + lTrack->GetDCAz() );
        StPhysicalHelixD helixGlob( vMomGlob, vOrigGlob,
                                   femtoEvent->GetMagneticField()*kilogauss,
                                   lTrack->GetCharge() );

        hdEdx->Fill( lTrack->GetDedx()*1e6 );
        hdEdx_p->Fill( lTrack->GetCharge()*lTrack->GetPtot(), lTrack->GetDedx()*1e6 );
        hP->Fill(lTrack->GetPtot());
        hPt->Fill(lTrack->GetPt());
        hDCAxy->Fill( lTrack->GetDCAxy() );
        hDCAz->Fill( lTrack->GetDCAz() );
        hDCAxyGlob->Fill( lTrack->GetDCAxyGlobal() );
        hDCAzGlob->Fill( lTrack->GetDCAzGlobal() );
        hEta->Fill( lTrack->GetEta() );

        if (lTrack->GetBeta() > 0.) {
          Float_t mInvBeta = 1./lTrack->GetBeta();

          hMSqr_P->Fill( lTrack->GetCharge()*lTrack->GetPtot(), lTrack->GetMassSqr() );
          hInvBeta_P->Fill( lTrack->GetCharge()*lTrack->GetPtot(), mInvBeta );
          hBeta_P->Fill( lTrack->GetCharge()*lTrack->GetPtot(), lTrack->GetBeta() );
        }

      }
    }
  }
  return kStOk;
}

  
void StFemtoDstQAMaker::SetTrigger(unsigned int trig, char mode)
{
  unsigned int trigColSize = mTrigIdCol.size();
  for (unsigned int iTrig = 0; iTrig < trigColSize; iTrig++) {
    if (mTrigIdCol[iTrig] == trig) {
      unsigned int mask = 1 << iTrig;
      if (!mode)
        mTrig = mTrig & (!mask);
      else
        mTrig = mTrig | mask;
      break;
    }
  }

  //
  // Output trigger mask
  //
  char bitMask[33] = {0};
  for (unsigned int iBit = 0; iBit < 32; iBit++) {
    bitMask[31 - iBit] = '0' + ((mTrig & (1 << iBit)) >> iBit);
  }
  std::cout << "StFemtoDstQAMaker[INFO]: Trigger mask = "
            << mTrig << " bit = "
            << bitMask << std::endl;
}
