#include "StPicoDstQA.h"


ClassImp(StPicoDstQA)


StPicoDstQA::StPicoDstQA(StPicoDstMaker *picoMaker, const char *outFilePath)
{
  std::cout << "Constructor" << std::endl;
  mPicoMaker = picoMaker;
  mOutFilePath = outFilePath;
}



StPicoDstQA::~StPicoDstQA()
{
  //
  // NO-OP
  //
}



Int_t StPicoDstQA::Init()
{
  std::cout << "Init" << std::endl;
  mOutFile = new TFile(mOutFilePath, "RECREATE");

  if (!mOutFile)
    std::cout << "Can't recreate file " << mOutFilePath << std::endl;

  Int_t mRefMultBins = 700;
  Double_t mRefMultLo = 0.;
  Double_t mRefMultHi = 700.;
  Int_t mPrimVtxZBins = 400;
  Double_t mPrimVtxZLo = -200.;
  Double_t mPrimVtxZHi = 200.;
  Int_t mPrimVertRBins = 60;
  Double_t mPrimVertRLo = -0.5;
  Double_t mPrimVertRHi = 1.0;
  Int_t mPrimVertXBins = 100;
  Double_t mPrimVertXLo = -1.;
  Double_t mPrimVertXHi = 1.;
  Int_t mPrimVertYBins = 100;
  Double_t mPrimVertYLo = -1.;
  Double_t mPrimVertYHi = 1.;
  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 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    mTofMultBins = 700;
  Double_t mTofMultLo = 0.;
  Double_t mTofMultHi = 700.;

  //
  // Event histograms
  //
  hRefMult = new TH1D("hRefMult", "reference multiplicity;RefMult;#",
                      mRefMultBins, mRefMultLo, mRefMultHi);
  hRefMult2 = new TH1D("hRefMult2", "reference multiplicity |#eta| < 1.;RefMult2;#",
                      mRefMultBins, mRefMultLo, mRefMultHi);
  hNGlobal = new TH1D("hNGlobal", "number of global tracks per event;N_{global};#",
                      1000, 0., 1000.);
  hNTracks = new TH1D("hNTracks", "number of tracks per event;nTracks;#",
                      300, 0., 300.);
  hNTracksDiff = new TH1D("hNTracksDiff", ";nTracksDiff;#",
                      600, -300., 300.);
  hR = new TH1D("hR", "R=#sqrt{v_{x}^{2} + v_{y}^{2}};R (cm);#",
                mPrimVertRBins, mPrimVertRLo, mPrimVertRHi);
  hVz = new TH1D("hVz", ";v_{z} (cm);#",
                 mPrimVtxZBins, mPrimVtxZLo, mPrimVtxZHi);
  hVxVy = new TH2D("hVxVy", ";V_{x};V_{y}",
                   mPrimVertXBins, mPrimVertXLo, mPrimVertXHi,
                   mPrimVertYBins, mPrimVertYLo, mPrimVertYHi);
  hVpdVzDiff = new TH1D("hVpdVzDiff", ";VpdV_{z} - V_{z};#",
                        300, -15., 15.);
  hTofMult = new TH1D("hTofMult", ";TofMult;#",
                      mTofMultBins, mTofMultLo, mTofMultHi);
  hRefMultForward = new TH1D("hRefMultForward ", "reference multiplicity |#eta| > 0.5;RefMultForward;#",
                      mRefMultBins, mRefMultLo, mRefMultHi);
  hTofVsRef = new TH2D("hTofVsRef", ";TofMult;RefMult",
                       mTofMultBins, mTofMultLo, mTofMultHi,
                       mRefMultBins, mRefMultLo, mRefMultHi);
  hTofVsRef2 = new TH2D("hTofVsRef2", ";TofMult;RefMult2",
                        mTofMultBins, mTofMultLo, mTofMultHi,
                        mRefMultBins, mRefMultLo, mRefMultHi);
  hRefVsRef2 = new TH2D("hRefVsRef2", ";RefMult;RefMult2",
                        mRefMultBins, mRefMultLo, mRefMultHi,
                        mRefMultBins, mRefMultLo, mRefMultHi);
  //
  // Run number dependence
  //
  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 TH2D("hNSigmaPionVsMom", "N#sigma(#pi) vs momenta;p (GeV/c);N#sigma(pi)",
                              mMomBins,mMomLo,mMomHi,
                              mNSigmaBins, mNSigmaLo, mNSigmaHi);
  hNSigmaKaonVsMom = new TH2D("hNSigmaKaonVsMom", "N#sigma(K) vs momenta;p (GeV/c);N#sigma(K)",
                              mMomBins,mMomLo,mMomHi,
                              mNSigmaBins, mNSigmaLo, mNSigmaHi);
  hNSigmaProtonVsMom = new TH2D("hNSigmaProtonVsMom", "N#sigma(p) vs momenta;p (GeV/c);N#sigma(p)",
                                mMomBins,mMomLo,mMomHi,
                                mNSigmaBins, mNSigmaLo, mNSigmaHi);
  hNSigmaElectronVsMom = new TH2D("hNSigmaElectronVsMom", "N#sigma(e) vs momenta;p (GeV/c);N#sigma(e)",
                                  mMomBins,mMomLo,mMomHi,
                                  mNSigmaBins, mNSigmaLo, mNSigmaHi);
  hP = new TH1D("hP", ";p (GeV/c);#", mMomBins, mMomLo, mMomHi);
  hGlobP = new TH1D("hGlobP", ";p (GeV/c);#", mMomBins, mMomLo, mMomHi);
  hEta = new TH1D("hEta", ";#eta;#",
                  200, -2., 2.);
  hPhi = new TH1D("hPhi", ";#phi;#",
                  400, -4., 4.);
  hEtaVsPhi = new TH2D("hEtaVsPhi",
                       ";#phi;#eta",
                       400, -4., 4.,
                       200, -2., 2.);
  hDedxVsPt = new TH2D("hDedxVsPt", ";p_{t}Q (GeV/c);dE/dx (keV/cm)",
                       2*mMomBins, -mMomHi, mMomHi,
                       mDedxBins, mDedxLo, mDedxHi);
  hMSqrVsPt = new TH2D("hMSqrVsPt",
                       ";p_{t}Q (GeV/c);m^{2} (GeV/c^{2})",
                       2*mMomBins, -mMomHi, mMomHi,
                       mMSqrBins, mMSqrLo, mMSqrHi);
  hInvBetaVsP = new TH2D("hInvBetaVsP",
                         "1/#beta vs p*q;Q*P (GeV/c);1/#beta",
                         2*mMomBins, -mMomHi, mMomHi, 
                         mInvBetaBins, mInvBetaLo, mInvBetaHi);
  hBetaVsP = new TH2D("hBetaVsP",
                      "#beta vs momentum;Q*P (GeV/c);#beta",
                      2*mMomBins, -mMomHi, mMomHi,
                      mBetaBins, mBetaLo, mBetaHi);
  hDcaXY = new TH1D("hDcaXY",
                    "Distance of closest approach (XY plane) for primary track;DCA_{xy} (cm);#",
                    mDCAxyBins, mDCAxyLo, mDCAxyHi);
  hDcaZ = new TH1D("hDcaZ",
                   "Distance of closest approach (Z axis) for primary track;DCA_{z} (cm);#",
                   mDCAzBins, mDCAzLo, mDCAzHi);
  hDca = new TH1D("hDca", "Distance of closest approach;DCA (cm);#", 100, 0., 50.);
  hDcaHelix = new TH1D("hDcaHelix", "", 100, 0., 50.);

  return StMaker::Init();
}


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


Int_t StPicoDstQA::Make()
{
  nEventsProc++;

  StPicoDst *picoDst = mPicoMaker->picoDst();
  if (!picoDst)
  {
    std::cout << "can't obtain StPicoDst" << std::endl;
    return kStOk;
  }

  StPicoEvent *picoEvent = picoDst->event();
  if (!picoEvent)
  {
    std::cout << "can't obtain StPicoEvent" << std::endl;
    return kStOk;
  }

  unsigned int trig = (unsigned int)picoEvent->triggerWord();
  if (trig & (!mTrigWord))
    return kStOk;

  unsigned int   nTracks    = picoDst->numberOfTracks();
  unsigned int   tofMult    = picoEvent->btofTrayMultiplicity();
  unsigned int   refMult    = picoEvent->refMult();
  unsigned int   refMultFrw = picoEvent->refMult2();
  unsigned int   refMult2   = refMult + refMultFrw;
  unsigned short globTracks = picoEvent->numberOfGlobalTracks();

  //
  // Fill event histograms
  //
  StThreeVectorF vtxPrim = picoEvent->primaryVertex();

  hRefMult->Fill(refMult);
  hRefMult2->Fill(refMult2);
  hRefMultForward->Fill(refMultFrw);
  hTofMult->Fill(tofMult);
  hTofVsRef->Fill(tofMult, refMult);
  hTofVsRef2->Fill(tofMult, refMult2);
  hRefVsRef2->Fill(refMult, refMult2);
  hNGlobal->Fill(globTracks);
  hNTracks->Fill(nTracks);
  hNTracksDiff->Fill(globTracks - nTracks);
  hR->Fill(vtxPrim.perp());
  hVz->Fill(vtxPrim.z());
  hVxVy->Fill(vtxPrim.x(),
              vtxPrim.y());
  hVpdVzDiff->Fill(picoEvent->vzVpd() - vtxPrim.z());

  //
  // Run number dependence
  //
  int runId = picoEvent->runId();

  hRefMultVsRunNumer->Fill(runId, refMult);
  hPrimTrackNumVsRunNumber->Fill(runId, nTracks);
  hGlobTrackNumVsRunNumber->Fill(runId, globTracks);
  hTofMultVsRunNumber->Fill(runId, tofMult);

  //
  // Loop over tracks
  //
  for (unsigned int iTrack = 0; iTrack < nTracks; iTrack++)
  {
    StPicoTrack *picoTrack = picoDst->track(iTrack);
    float p      = picoTrack->pMom().mag();
    if (p < 0.1) continue;
    float pGlob  = picoTrack->gMom().mag();
    float pt     = picoTrack->pMom().perp();
    float eta    = picoTrack->pMom().pseudoRapidity();
    float dedx   = picoTrack->dEdx();
    short charge = picoTrack->charge();
    float phi    = picoTrack->pMom().phi();
    //
    // Get beta
    //
    float beta = picoTrack->btofBeta();
    if (beta >= 0.)
    {
      float invBeta = 1./beta;
      float massSqr = p*p*(invBeta*invBeta - 1.);
      hMSqrVsPt->Fill(pt*charge, massSqr);
      hBetaVsP->Fill(beta, p);
      hInvBetaVsP->Fill(invBeta, p);
    }
    //
    // Calculate beta
    //
    StPhysicalHelixD globHelix(picoTrack->gMom(),
                               picoTrack->origin(),
                               picoEvent->bField() * kilogauss,
                               charge);
    double pathLen = globHelix.pathLength(vtxPrim, false);
    StThreeVectorF dca = globHelix.at(pathLen) - vtxPrim;

    hNSigmaPionVsMom->Fill(p, picoTrack->nSigmaPion());
    hNSigmaKaonVsMom->Fill(p, picoTrack->nSigmaKaon());
    hNSigmaProtonVsMom->Fill(p, picoTrack->nSigmaProton());
    hNSigmaElectronVsMom->Fill(p, picoTrack->nSigmaElectron());

    hP->Fill(p);
    hEta->Fill(eta);
    hPhi->Fill(phi);
    hEtaVsPhi->Fill(phi, eta);
    hGlobP->Fill(pGlob);
    hDedxVsPt->Fill(pt*charge, dedx);
    hDca->Fill(picoTrack->dca());
    hDcaHelix->Fill(dca.mag());
    hDcaXY->Fill(dca.perp());
    hDcaZ->Fill(dca.z());
  }

  return kStOk;
}


Int_t StPicoDstQA::Finish()
{
  std::cout << "Finish!" << std::cout;
  mOutFile->cd();
  mOutFile->Write();
  mOutFile->Close();

  return kStOk;
}
