#include <TROOT.h>
#include <TSystem.h>
#include <TChain.h>
#include <TFile.h>
#include <TString.h>
#include <iostream>

//
// Constants for analyses
//

#define MULT_BINS 3
#define CHARGE_BINS 2

//
//  Event information
//
const int mPrimVertZbins = 6;                // number of z bins for analysis
const double mPrimVertZ[2] = {-30., 30.};    // redefine values of the vertex Z for each run [cm]
const double mPrimVertR = 2.;                // distance in the XY plane [cm]
const double mPrimVertVpdVz[2] = {-5., 5.};  // VpdZ - TPCZ difference
const double mPrimVertXShift = 0.;
const double mPrimVertYShift = 0.;
const bool mUseZdcCorrection = false;         // valid for AuAu collisions. True for 200GeV only!

//
// Track information
//
const double mChargedKaonMass = 0.493677;
const double mChargedPionMass = 0.13957018;
const int mPosCharge = +1;
const int mNegCharge = -1;
const int mTrackType = 1;                     // primary track
const int mPythiaKaonCode = 321;              // pythia particle code
const int mPythiaPionCode = 211;              // pythia particle code

const int mTrackNHits[2] = {15, 50};          // number of track hits
const double mTrackNHitsFit[2] = {5, 50};    // number of track hits used in fit
const double mTrackDCA[2] = {0., 2.};         // DCA of the track and PV [cm]
const double mTrackASplitVal = 0.;
const double mTrackMom[2] = {0.15, 1.55};        // general momentum of the track [GeV/c]
const double mTrackTransvMom[2] = {0.05, 1.55};  // general transverse momentum of the track [GeV/c]
const double mTrackEta[2] = { -1., 1.};

const short mDetectorSelection = 4;  // 0-TPC, 1-ToF, 2-TPC+ToF, 3-ToF||TPC, 4-TPC&&TOF(P>Pthr)||TPC(P<Pthr)

const double mTrackTpcMom[2] = {0.15, 0.55};  // track momentum for TPC idnetification [GeV/c]
const double mnSigmaElectron[2] = {-1e9, 1e9};
const double mnSigmaPion[2] = {-1e9, 1e9};
const double mnSigmaKaon[2] = {-2., 2.};
const double mnSigmaProton[2] = {-1e9, 1e9};

const double mTrackTofMom[2] = {0.15, 1.45};   // track momentum for ToF identification [GeV/c]
const double mPionTofMsqr[2] = {-0.02, 0.062}; // based on p+p data
const double mKaonTofMsqr[2] = {0.2, 0.32};    // based on p+p data

const double mTrackTnTMom[2] = {0.15, 1.55};
const double mTnTNSigmaElectron[2] = {-1e9, 1e9};
const double mTnTNSigmaPion[2] = {-1e9, 1e9};
const double mTnTNSigmaKaon[2] = {-2., 2.};
const double mTnTNSigmaProton[2] = {-1e9, 1e9};

const double mTTTMomThresh = 0.5; // TPC&&TOF (P>Pthr)||TPC(P<Pthr)

//
// Pair information
//
const double mFmrCut[2] = {-1.1, 0.1};         // not more than 10% of merged rows
const double mSplitLevelCut[2] = {-0.5, 0.6};  // splitting level from paper. StHbtPair::quality()
const double mPairKt[2] = {0.,1.5};            // pair transverse momentum cut
const double mPairQinv[2] = {0, 2.4};          // pair relative momentum
const double mAveSep[2] = {5., 5000.};         // average separation of tracks within TPC

//
// Histograms values
//
const int mPairKtBins = 10;                 // bin width
const double mPairKtVal[2] = {0.1, 1.1};    // 100 MeV/c
const int mQinvBins = 120;                   // 20 MeV/c bin width
const double mQinvVal[2] = {0., 2.4};        // redefine for each analysis

const Char_t *defaultInFile = "kaons.list";
const Char_t *defaultOutFile = "oExpKaon_pp200_1D.root";

//_________________
void expKaonHBT_pp200_PicoDst_1D_v2(const Char_t *inFileList = defaultInFile,
                                    const Char_t *oFileName  = defaultOutFile) {


  std::cout << "********************************" << std::endl
      << "*       expKaonHBT_pp200_1D    *" << std::endl
      << "*               Start          *" << std::endl
      << "********************************" << std::endl << endl;

  //
  // Load libraries
  //
  std::cout << "Loading libraries..." << std::endl;
  gROOT->LoadMacro("$STAR/StRoot/StMuDSTMaker/COMMON/macros/loadSharedLibraries.C");
  loadSharedLibraries();
  gSystem->Load("libMinuit");
  gSystem->Load("StRefMultCorr");
  gSystem->Load("StFlowMaker");   // should be placed before HBT
  gSystem->Load("StHbtMaker");
  gSystem->Load("StarClassLibrary");
  gSystem->Load("libgsl");
  gSystem->Load("libgslcblas");
  gSystem->Load("StPicoDstMaker");

  std::cout << "Libraries have been successfully loaded" << std::endl;

  //
  // Create chain
  //
  StChain* mChain = new StChain("StChain");
  mChain->SetDebug(0);
  StMuDebug::setLevel(0);

  //
  // Create PicoDst maker
  //
  StPicoDstMaker *picoMaker = new StPicoDstMaker(0, inFileList);

  //
  // Set StHbtMaker and StHbtManager
  //
  StHbtMaker *mHbtMaker = new StHbtMaker();
  mHbtMaker->SetDebug(0);
  StHbtManager *mHbtManager = mHbtMaker->HbtManager();

  //
  // Set StHbtReader
  //
  StPicoDstHbtReader *picoReader = new StPicoDstHbtReader(picoMaker);
  unsigned int trigWord = 0;
  for (int i = 0; i < 2; trigWord |= 1 << i++); // set only VPDMB
  std::cout << "TRIGGER WORD: " << trigWord << std::endl;
  picoReader->SetTriggers(trigWord);

  mHbtManager->SetEventReader(picoReader);

  //
  // Set kaon pair cut
  //
  gregsTrackPairCut *mKaonPairCut = new gregsTrackPairCut();
  mKaonPairCut->SetFracOfMergedRow(mFmrCut[0],mFmrCut[1]);
  mKaonPairCut->SetQuality(mSplitLevelCut[0],mSplitLevelCut[1]);
  mKaonPairCut->SetKt(mPairKt[0],mPairKt[1]);
  mKaonPairCut->SetQinv(mPairQinv[0],mPairQinv[1]);
  mKaonPairCut->SetAverageSeparation(mAveSep[0],mAveSep[1]);

  Int_t mLocalCharge, mMultBins, mMultLow, mMultHi, mEvents2Mix, mCurrentPid;

  //
  // The multiplicity analysis matrix is CHARGE_BINS x MULT_BINS, where
  // the first element corresponds to the particle charge: (pos, neg),
  // and the second one corresponds to the RefMult intervals as listed below:
  // Mult: 0-45, 0-2, 3-4, 5-6, 7-8, 9-10, 11-45
  //

  StHbtVertexMultAnalysis *mKaonMultAnalysis[CHARGE_BINS][MULT_BINS]; // StHbtAnalysis, cut, CF and monitor declarations

  //
  // Cuts
  //
  gregsEventCut       *mKaonEventCut[MULT_BINS];
  gregsSimpleTrackCut *mKaonTrackCut[CHARGE_BINS][MULT_BINS];

  //
  // CorrFctns
  //
  QinvCorrFctnKt                *mKaonQinvMixKtCF[CHARGE_BINS];
  QinvCorrFctnOpposHemisphereKt *mKaonQinvOppKtCF[CHARGE_BINS];
  QinvCorrFctnRotationKt        *mKaonQinvRotKtCF[CHARGE_BINS];
  gregsBPLCMSFrame3DCorrFctnKt  *mKaonQinvBPlcmsKtCF[CHARGE_BINS];
  QinvCorrFctnKt_MomCons        *mKaonQinvMomConsKtCF[CHARGE_BINS];
  QinvCorrFctn                  *mKaonQinvMixMultCF[CHARGE_BINS][MULT_BINS];
  QinvCorrFctnOpposHemisphere   *mKaonQinvOppMultCF[CHARGE_BINS][MULT_BINS];
  QinvCorrFctnRotation          *mKaonQinvRotMultCF[CHARGE_BINS][MULT_BINS];
  gregsBPLCMSFrame3DCorrFctn    *mKaonQinvBlcmsCF[CHARGE_BINS][MULT_BINS];
  QinvCorrFctn_MomCons          *mKaonQinvMomConsCF[CHARGE_BINS][MULT_BINS];

  //
  // Analysis creation
  //
  // Multiplicity loop
  for(Int_t iMult = 0; iMult < MULT_BINS; iMult++) {
    //
    // Setup event cut
    //
    switch(iMult) {
      case 1:  mMultLow = 0;  mMultHi = 7;   mMultBins = 1;  mEvents2Mix = 10; break;
      case 2:  mMultLow = 8;  mMultHi = 70;  mMultBins = 1;  mEvents2Mix = 10; break;
      default: mMultLow = 0;  mMultHi = 70;  mMultBins = 3;  mEvents2Mix = 10; break;
    };

    mKaonEventCut[iMult] = new gregsEventCut();
    mKaonEventCut[iMult]->SetCollisionType(0);     //0-pp, 1-AuAu
    mKaonEventCut[iMult]->SetVertZPos(mPrimVertZ[0], mPrimVertZ[1]);
    mKaonEventCut[iMult]->SetVzVpdVzDiff(mPrimVertVpdVz[0],mPrimVertVpdVz[1]);
    mKaonEventCut[iMult]->SetVertRPos(mPrimVertR);
    mKaonEventCut[iMult]->SetEventMult(mMultLow, mMultHi);

    //
    // Loop over charge of the particles
    //
    for(Int_t iCharge = 0; iCharge < CHARGE_BINS; iCharge++) {
      switch(iCharge) {
        case 0:   mLocalCharge = mPosCharge;    break;
        default:  mLocalCharge = mNegCharge;    break;
      };

      //
      // Like-sign kaon analysis
      //
      // Create analysis
      mKaonMultAnalysis[iCharge][iMult] = new StHbtVertexMultAnalysis(mPrimVertZbins,
                                                                      mPrimVertZ[0],
                                                                      mPrimVertZ[1],
                                                                      mMultBins,
                                                                      mMultLow,
                                                                      mMultHi);

      mKaonMultAnalysis[iCharge][iMult]->SetEventCut(mKaonEventCut[iMult]);

      //
      // Create single-track cuts and add them to the analysis
      //
      mKaonTrackCut[iCharge][iMult] = new gregsSimpleTrackCut();
      mKaonTrackCut[iCharge][iMult]->SetCharge(mLocalCharge);
      mKaonTrackCut[iCharge][iMult]->SetMass(mChargedKaonMass);
      mKaonTrackCut[iCharge][iMult]->SetType(mTrackType);
      mKaonTrackCut[iCharge][iMult]->SetNHits(mTrackNHits[0], mTrackNHits[1]);
      mKaonTrackCut[iCharge][iMult]->SetNHitsFit(mTrackNHitsFit[0], mTrackNHitsFit[1]);
      mKaonTrackCut[iCharge][iMult]->SetAntiSplit(mTrackASplitVal);
      mKaonTrackCut[iCharge][iMult]->SetEta(mTrackEta[0], mTrackEta[1]);
      mKaonTrackCut[iCharge][iMult]->SetDCA(mTrackDCA[0], mTrackDCA[1]);
      mKaonTrackCut[iCharge][iMult]->SetDCAGlobal(mTrackDCA[0], mTrackDCA[1]);
      mKaonTrackCut[iCharge][iMult]->SetP(mTrackMom[0], mTrackMom[1]);

      mKaonTrackCut[iCharge][iMult]->SetDetectorSelection(mDetectorSelection);

      mKaonTrackCut[iCharge][iMult]->SetTpcMomentum(mTrackTpcMom[0], mTrackTpcMom[1]);
      mKaonTrackCut[iCharge][iMult]->SetNSigmaElectron(mnSigmaElectron[0],mnSigmaElectron[1]);
      mKaonTrackCut[iCharge][iMult]->SetNSigmaPion(mnSigmaPion[0], mnSigmaPion[1]);
      mKaonTrackCut[iCharge][iMult]->SetNSigmaKaon(mnSigmaKaon[0], mnSigmaKaon[1]);
      mKaonTrackCut[iCharge][iMult]->SetNSigmaProton(mnSigmaProton[0], mnSigmaProton[1]);

      mKaonTrackCut[iCharge][iMult]->SetTofMassSqr(mKaonTofMsqr[0],mKaonTofMsqr[1]);
      mKaonTrackCut[iCharge][iMult]->SetTofMomentum(mTrackTofMom[0], mTrackTofMom[1]);

      mKaonTrackCut[iCharge][iMult]->SetTnTMomentum(mTrackTnTMom[0], mTrackTnTMom[1]);
      mKaonTrackCut[iCharge][iMult]->SetTnTNSigmaElectron(mTnTNSigmaElectron[0], mTnTNSigmaElectron[1]);
      mKaonTrackCut[iCharge][iMult]->SetTnTNSigmaPion(mTnTNSigmaPion[0], mTnTNSigmaPion[1]);
      mKaonTrackCut[iCharge][iMult]->SetTnTNSigmaKaon(mTnTNSigmaKaon[0], mTnTNSigmaKaon[1]);
      mKaonTrackCut[iCharge][iMult]->SetTnTNSigmaProton(mTnTNSigmaProton[0], mTnTNSigmaProton[1]);

      mKaonTrackCut[iCharge][iMult]->SetTTTMomentumThresh(mTTTMomThresh);

      mKaonMultAnalysis[iCharge][iMult]->SetFirstParticleCut(mKaonTrackCut[iCharge][iMult]);
      mKaonMultAnalysis[iCharge][iMult]->SetSecondParticleCut(mKaonTrackCut[iCharge][iMult]);

      //
      // Create pair cut
      //
      mKaonMultAnalysis[iCharge][iMult]->SetPairCut(mKaonPairCut);

      //
      // Correlation function initialization
      //
      if (iMult == 0) {
        mKaonQinvMixKtCF[iCharge] = 
            new QinvCorrFctnKt(Form("hKaonQinvMixKt_%d", iCharge),
                               mQinvBins, mQinvVal[0], mQinvVal[1],
                               mPairKtBins, mPairKtVal[0], mPairKtVal[1]);
        mKaonQinvOppKtCF[iCharge] = 
            new QinvCorrFctnOpposHemisphereKt(Form("hKaonQinvOppKt_%d", iCharge),
                                              mQinvBins, mQinvVal[0], mQinvVal[1],
                                              mPairKtBins, mPairKtVal[0], mPairKtVal[1]);
        mKaonQinvRotKtCF[iCharge] = 
            new QinvCorrFctnRotationKt(Form("hKaonQinvRotKt_%d", iCharge),
                                       mQinvBins, mQinvVal[0], mQinvVal[1],
                                       mPairKtBins, mPairKtVal[0], mPairKtVal[1]);
        mKaonQinvBPlcmsKtCF[iCharge] = 
            new gregsBPLCMSFrame3DCorrFctnKt(Form("hKaonQinvBPlcmsKtCF_%d", iCharge),
                                             mQinvBins, mQinvVal[0], mQinvVal[1],
                                             mPairKtBins, mPairKtVal[0], mPairKtVal[1]);
        mKaonQinvMomConsKtCF[iCharge] =
            new QinvCorrFctnKt_MomCons(Form("hKaonQinvMomConsKtCF_%d", iCharge),
                                       mQinvBins, mQinvVal[0], mQinvVal[1],
                                       mPairKtBins, mPairKtVal[0], mPairKtVal[1]);

        mKaonMultAnalysis[iCharge][iMult]->AddCorrFctn(mKaonQinvMixKtCF[iCharge]);
        mKaonMultAnalysis[iCharge][iMult]->AddCorrFctn(mKaonQinvOppKtCF[iCharge]);
        mKaonMultAnalysis[iCharge][iMult]->AddCorrFctn(mKaonQinvRotKtCF[iCharge]);
        mKaonMultAnalysis[iCharge][iMult]->AddCorrFctn(mKaonQinvBPlcmsKtCF[iCharge]);
        mKaonMultAnalysis[iCharge][iMult]->AddCorrFctn(mKaonQinvMomConsKtCF[iCharge]);
      }
      mKaonQinvMixMultCF[iCharge][iMult] = 
          new QinvCorrFctn(Form("hKaonQinvMixMult_%d_%d", iCharge, iMult),
                           mQinvBins, mQinvVal[0], mQinvVal[1]);
      mKaonQinvOppMultCF[iCharge][iMult] = 
          new QinvCorrFctnOpposHemisphere(Form("hKaonQinvOppMult_%d_%d", iCharge, iMult),
                                          mQinvBins, mQinvVal[0], mQinvVal[1]);
      mKaonQinvRotMultCF[iCharge][iMult] = 
          new QinvCorrFctnRotation(Form("hKaonQinvRotMult_%d_%d", iCharge, iMult),
                                   mQinvBins, mQinvVal[0], mQinvVal[1]);
      mKaonQinvBlcmsCF[iCharge][iMult] = 
          new gregsBPLCMSFrame3DCorrFctn(Form("hKaonQinvBlcmsCF_%d_%d", iCharge, iMult),
                                         mQinvBins, mQinvVal[0], mQinvVal[1]);
      mKaonQinvMomConsCF[iCharge][iMult] =
          new QinvCorrFctn_MomCons(Form("hKaonQinvMomConsCF_%d_%d", iCharge, iMult),
                                   mQinvBins, mQinvVal[0], mQinvVal[1]);

      mKaonMultAnalysis[iCharge][iMult]->AddCorrFctn(mKaonQinvMixMultCF[iCharge][iMult]);
      mKaonMultAnalysis[iCharge][iMult]->AddCorrFctn(mKaonQinvOppMultCF[iCharge][iMult]);
      mKaonMultAnalysis[iCharge][iMult]->AddCorrFctn(mKaonQinvRotMultCF[iCharge][iMult]);
      mKaonMultAnalysis[iCharge][iMult]->AddCorrFctn(mKaonQinvBlcmsCF[iCharge][iMult]);
      mKaonMultAnalysis[iCharge][iMult]->AddCorrFctn(mKaonQinvMomConsCF[iCharge][iMult]);

      mKaonMultAnalysis[iCharge][iMult]->SetNumEventsToMix(mEvents2Mix); // set number of events to mix
      mHbtManager->AddAnalysis(mKaonMultAnalysis[iCharge][iMult]); // add analysis to the manager
    }
  }

  mChain->Init(); // StChain initialization
  Long64_t mNEvents = picoMaker->chain()->GetEntries();

  //
  // Start processing
  //
  Int_t iRet = 0;
  Int_t mNEventsProcessed = 0;
  Float_t mPercentCounter = 0.05;
  Float_t mProgress = 0.;

  unsigned int mSubCounter = 1;              // time
  time_t mStartTime, mStopTime, mDiffTime;   // time
  float mFrac;                               // time
  mStartTime = time(0);                      // time

  // mNEvents = 15;
  std::cout << "Total Events: " << mNEvents << std::endl;  
  for(Int_t iEvent=0; iEvent < mNEvents; iEvent++) {

    mProgress = (Float_t)iEvent / (Float_t)mNEvents;
    mNEventsProcessed++;

    if(mProgress >= mPercentCounter) {
      mPercentCounter += 0.05;
      mStopTime = time(0);
      mDiffTime = difftime(mStopTime, mStartTime);
      mFrac = (float)mDiffTime*(float)(mNEvents - iEvent)/(float)iEvent;
      mSubCounter++;
      std::cout << Form("Processing progress: %4.2f%%. Time left: %.1f sec", mProgress*100., mFrac)
          << std::endl;
    }

    mChain->Clear();
    iRet = mChain->Make(iEvent);
    if(iRet!=0) {
      std::cout << "expMesonHBT_pp200_1D[ERROR] - iRet: " << iRet << std::endl;
      break;
    }
  }

  mChain->Finish();

  TFile *oFile = new TFile(oFileName, "RECREATE"); // create output file

  //
  // Write hbt histograms to the file
  //
  std::cout << "Writing histograms to the output file...";

  for(Int_t iMult = 0; iMult < MULT_BINS; iMult++)
  {
    for(Int_t iCharge = 0; iCharge < CHARGE_BINS; iCharge++)
    {
      if (iMult == 0)
      {
        for(Int_t iKt = 0; iKt < mPairKtBins; iKt++)
        {
          mKaonQinvMixKtCF[iCharge]->Numerator(iKt)->Write();
          mKaonQinvMixKtCF[iCharge]->Denominator(iKt)->Write();
          mKaonQinvOppKtCF[iCharge]->Numerator(iKt)->Write();
          mKaonQinvOppKtCF[iCharge]->Denominator(iKt)->Write();
          mKaonQinvRotKtCF[iCharge]->Numerator(iKt)->Write();
          mKaonQinvRotKtCF[iCharge]->Denominator(iKt)->Write();
          mKaonQinvBPlcmsKtCF[iCharge]->Numerator(iKt)->Write();
          mKaonQinvBPlcmsKtCF[iCharge]->Denominator(iKt)->Write();
          mKaonQinvMomConsKtCF[iCharge]->Numerator(iKt)->Write();
          mKaonQinvMomConsKtCF[iCharge]->Denominator(iKt)->Write();
        }
      }
      mKaonQinvMixMultCF[iCharge][iMult]->Numerator()->Write();
      mKaonQinvMixMultCF[iCharge][iMult]->Denominator()->Write();
      mKaonQinvOppMultCF[iCharge][iMult]->Numerator()->Write();
      mKaonQinvOppMultCF[iCharge][iMult]->Denominator()->Write();
      mKaonQinvRotMultCF[iCharge][iMult]->Numerator()->Write();
      mKaonQinvRotMultCF[iCharge][iMult]->Denominator()->Write();
      mKaonQinvBlcmsCF[iCharge][iMult]->Numerator()->Write();
      mKaonQinvBlcmsCF[iCharge][iMult]->Denominator()->Write();
      mKaonQinvMomConsCF[iCharge][iMult]->Numerator()->Write();
      mKaonQinvMomConsCF[iCharge][iMult]->Denominator()->Write();
    }
  }

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

  //
  // Write the output file and close it
  //
  std::cout << "Saving the output file...";
  oFile->Close();
  std::cout << "\t[DONE]" << std::endl;

  delete mChain;

  std::cout << "********************************" << std::endl
      << "*       expKaonHBT_pp200_1D    *" << std::endl
      << "*              Finish          *" << std::endl
      << "********************************" << std::endl << std::endl;
}
