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

// #define MONITORS

//Constants for analyses
//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_t mUseZdcCorrection = false;

//Track information
const double mChargedKaonMass = 0.493677;
const double mChargedPionMass = 0.13957018;
const int mPosCharge = +1;
const int mNegCharge = -1;
const int mTrackType = 1; 
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] = {15, 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.52;
const double mTrackMom[2] = {0.15, 1.55};     //General momentum of the track [GeV/c]
const double mTrackTransvMom[2] = {0.15, 1.55};  //General transverse momentum of the track [GeV/c]
const double mTrackEta[2] = { -1., 1.};

const short mDetectorSelection = 3;   

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]     = {-2.,  2.};
const double mnSigmaKaon[2]     = {-1e9, 1e9};
const double mnSigmaProton[2]   = {-1e9, 1e9};

const double mTrackTofMom[2] = {0.15, 1.55};   //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.20, 0.32};    //Based on p+p data

const double mTrackTnTMom[2]   = {0.15, 1.55};
const double mTnTNSigmaPion[2] = {-2., 2.};

//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., 1.5};        //Pair relative momentum
const double mAveSep[2] = {5., 5000.};         //Average separation of tracks within TPC

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

const Char_t *defaultInFile = "sample.list";
const Char_t *defaultOutFile = "oSimPion_pau200_3D.root";

//_________________
void expPionDCHBT_pau200_y2015_3D(const Char_t *inFileList = defaultInFile,
				  const Char_t *oFileName = defaultOutFile,
				  int mult = 0) {

  std::cout << "**************************************" << std::endl
	    << "*     simPionDCHBT_pau200_3D         *" << 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("StHbtStarGenHijingDstReader");
  gSystem->Load("StFemtoDstMaker");
  gSystem->Load("libgsl");
  gSystem->Load("libgslcblas");

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

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

  // set StHbtReader
  std::cout << Form("Try to read file: %s", inFileList) << std::endl;
  
  StHbtStarGenHijingDstReader *hijingReader = new StHbtStarGenHijingDstReader("", inFileList, "root", -1, false);
  Long64_t mNEvents = hijingReader->GetNEvents();
  
  std::cout << Form("Number of events in TTree: %d", mNEvents) << std::endl;

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

  // set basic event cut
  gregsFemtoDstEventCut *mBasicEventCut = new gregsFemtoDstEventCut();
  mBasicEventCut->SetCollisionType(0);     // 0-pp, 1-AuAu
  mBasicEventCut->SetEventMult(0, 70);
  mBasicEventCut->SetVertZPos(mPrimVertZ[0], mPrimVertZ[1]);
  mBasicEventCut->SetVzVpdVzDiff(mPrimVertVpdVz[0],mPrimVertVpdVz[1]);
  mBasicEventCut->SetVertRPos(mPrimVertR);
  mBasicEventCut->SetEventCentralityBins(0, 8);  // for 5 centrality bins

  // set basic pion single-particle cut
  gregsPionTrackCut *mBasicTrackCut = new gregsPionTrackCut();
  mBasicTrackCut->SetCharge(mPosCharge);
  mBasicTrackCut->SetMass(mChargedPionMass);
  mBasicTrackCut->SetType(mTrackType);
  mBasicTrackCut->SetNHits(mTrackNHits[0], mTrackNHits[1]);
  mBasicTrackCut->SetNHitsFit(mTrackNHitsFit[0], mTrackNHitsFit[1]);
  mBasicTrackCut->SetAntiSplit(mTrackASplitVal);
  mBasicTrackCut->SetEta(mTrackEta[0], mTrackEta[1]);
  mBasicTrackCut->SetDCA(mTrackDCA[0], mTrackDCA[1]);
  mBasicTrackCut->SetDCAGlobal(mTrackDCA[0], mTrackDCA[1]);
  mBasicTrackCut->SetP(mTrackMom[0], mTrackMom[1]);

  mBasicTrackCut->SetDetectorSelection(mDetectorSelection);

  mBasicTrackCut->SetTpcMomentum(mTrackTpcMom[0], mTrackTpcMom[1]);
  mBasicTrackCut->SetTpcNSigmaElectron(mnSigmaElectron[0],mnSigmaElectron[1]);
  mBasicTrackCut->SetTpcNSigmaPion(mnSigmaPion[0], mnSigmaPion[1]);
  mBasicTrackCut->SetTpcNSigmaKaon(mnSigmaKaon[0], mnSigmaKaon[1]);
  mBasicTrackCut->SetTpcNSigmaProton(mnSigmaProton[0], mnSigmaProton[1]);

  mBasicTrackCut->SetTofMassSqr(mPionTofMsqr[0],mPionTofMsqr[1]);
  mBasicTrackCut->SetTofMomentum(mTrackTofMom[0], mTrackTofMom[1]);

  mBasicTrackCut->SetTnTMomentum(mTrackTnTMom[0], mTrackTnTMom[1]);
  mBasicTrackCut->SetTnTNSigmaPion(mTnTNSigmaPion[0], mTnTNSigmaPion[1]);

  // set basic pion pair cut
  gregsTrackPairCut *mBasicPairCut = new gregsTrackPairCut();
  // mBasicPairCut->SetFracOfMergedRow(mFmrCut[0],mFmrCut[1]);
  // mBasicPairCut->SetQuality(mSplitLevelCut[0],mSplitLevelCut[1]);
  mBasicPairCut->SetKt(mPairKt[0],mPairKt[1]);
  mBasicPairCut->SetQinv(mPairQinv[0],mPairQinv[1]);
  // mBasicPairCut->SetAverageSeparation(mAveSep[0],mAveSep[1]);

  Int_t mLocalCharge, mMultBins, mMultLow, mMultHi, mEvents2Mix;
  float ktLow, ktHi;

  // The multiplicity analysis array, where
  // the first element corresponds to the particle charge: (pos, neg),
  // and 6 centralitiy bins: (1-100),(1-6),(7-9),(10-12),(13-16),(17,100)

  // StHbtAnalysis, cut, CF and monitor declarations
  StHbtVertexMultAnalysis *mPionMultAnalysis;
  
  // cuts
  gregsFemtoDstEventCut *mPionMultEventCut;
  gregsPionTrackCut     *mPionPosMultTrackCut, *mPionNegMultTrackCut;
  gregsTrackPairCut     *mPionMultPairCut;

  // correlation functions
  QinvCorrFctnKt               *mPionQinvKtCF;
  gregsBPLCMSFrame3DCorrFctnKt *mPionBPLCMSFrameKtCF;

#ifdef MONITORS
  //Cut monitors
  fxtEventCutMonitor *mFxtPionEventCutMonitorPass;
  fxtTrackCutMonitor *mFxtPionPosTrackCutMonitorPass, *mFxtPionNegTrackCutMonitorPass;
  fxtPairCutMonitor  *mFxtPionPairCutMonitorPass;
  fxtEventCutMonitor *mFxtPionEventCutMonitorFail;
  fxtTrackCutMonitor *mFxtPionTrackCutMonitorFail;
  fxtPairCutMonitor  *mFxtPionPairCutMonitorFail;
#endif

  std::cout << "---=== BEGIN CREATE ANALYSIS ===---\n";

  switch(mult) {
  case 1:  mMultLow = 0;  mMultHi = 6;    mMultBins = 2; mEvents2Mix = 5; break;
  case 2:  mMultLow = 7;  mMultHi = 9;    mMultBins = 2; mEvents2Mix = 5; break;
  case 3:  mMultLow = 10; mMultHi = 12;   mMultBins = 2; mEvents2Mix = 3; break;
  case 4:  mMultLow = 13; mMultHi = 16;   mMultBins = 2; mEvents2Mix = 3; break;
  case 5:  mMultLow = 17; mMultHi = 25;   mMultBins = 2; mEvents2Mix = 3; break;
  case 6:  mMultLow = 26; mMultHi = 35;   mMultBins = 2; mEvents2Mix = 3; break;
  case 7:  mMultLow = 36; mMultHi = 45;   mMultBins = 2; mEvents2Mix = 3; break;
  case 8:  mMultLow = 46; mMultHi = 70;   mMultBins = 2; mEvents2Mix = 3; break;      
  default: mMultLow = 0;  mMultHi = 70;   mMultBins = 7; mEvents2Mix = 3; break;
  };

  //Create analysis
  mPionMultAnalysis = new StHbtVertexMultAnalysis(mPrimVertZbins,
						  mPrimVertZ[0],
						  mPrimVertZ[1],
						  mMultBins,
						  mMultLow,
						  mMultHi);
    
  //Create event cut
  mPionMultEventCut = new gregsFemtoDstEventCut(*mBasicEventCut);
  mPionMultEventCut->SetEventMult(mMultLow, mMultHi);
	
#ifdef MONITORS
  mFxtPionEventCutMonitorPass = 
    new fxtEventCutMonitor("eventPionPass", Form("_%d", mult));
  mFxtPionEventCutMonitorFail = 
    new fxtEventCutMonitor("eventPionFail", Form("_%d", mult));
  mPionEventCut->AddCutMonitor(mFxtPionEventCutMonitorPass,
			       mFxtPionEventCutMonitorFail);
#endif
  mPionMultAnalysis->SetEventCut(mPionMultEventCut);

  //
  // Create single-track cuts and add them to the analysis
  //
  mPionPosMultTrackCut = new gregsPionTrackCut(*mBasicTrackCut);
  mPionPosMultTrackCut->SetCharge(mPosCharge);
  mPionNegMultTrackCut = new gregsPionTrackCut(*mBasicTrackCut);
  mPionNegMultTrackCut->SetCharge(mNegCharge);
    
#ifdef MONITORS
  //
  // Positive charged pions cut monitors
  //
  mFxtPionPosTrackCutMonitorPass = 
    new fxtTrackCutMonitor(Form("_trackPionPosPass_%d_", mult), mChargedPionMass);
  mFxtPionPosTrackCutMonitorFail = 
    new fxtTrackCutMonitor(Form("_trackPionPosFail_%d_", mult), mChargedPionMass);
  mPionPosMultTrackCut->AddCutMonitor(mFxtPionPosTrackCutMonitorPass,
				      mFxtPionPosTrackCutMonitorFail);
  //
  // Negative charged pions cut monitors
  //
  mFxtPionNegTrackCutMonitorPass = 
    new fxtTrackCutMonitor(Form("_trackPionNegPass_%d_", mult), mChargedPionMass);
  mFxtPionNegTrackCutMonitorFail = 
    new fxtTrackCutMonitor(Form("_trackPionNegFail_%d_", mult), mChargedPionMass);
  mPionNegMultTrackCut->AddCutMonitor(mFxtPionNegTrackCutMonitorPass,
				      mFxtPionNegTrackCutMonitorFail);
    
#endif
  mPionMultAnalysis->SetFirstParticleCut(mPionPosMultTrackCut);
  mPionMultAnalysis->SetSecondParticleCut(mPionNegMultTrackCut);

  //Create pair cut
  mPionMultPairCut = new gregsTrackPairCut(*mBasicPairCut);

#ifdef MONITORS
  mFxtPionPairCutMonitorPass = 
    new fxtPairCutMonitor(Form("_pairPionPass_%d_", mult));
  mFxtPionPairCutMonitorPass->SetParticleMass(mChargedPionMass);
  mFxtPionPairCutMonitorFail = 
    new fxtPairCutMonitor(Form("_pairPionFail_%d_", mult));
  mFxtPionPairCutMonitorFail->SetParticleMass(mChargedPionMass);
  mPionPairCut->AddCutMonitor(mFxtPionPairCutMonitorPass,
			      mFxtPionPairCutMonitorFail);
#endif
  mPionMultAnalysis->SetPairCut(mPionMultPairCut);
  //
  // Correlation function initialization
  //
  // 1D
  mPionQinvKtCF = 
    new QinvCorrFctnKt(Form("hPionDC_QinvKtMix_%d", mult),
		       mQinvBins, mQinvVal[0], mQinvVal[1],
		       mPairKtBins, mPairKtVal[0], mPairKtVal[1]);
  // 3D
  mPionBPLCMSFrameKtCF = 
    new gregsBPLCMSFrame3DCorrFctnKt(Form("hPionDC_BPLCMSFrameKtCF_%d", mult),
				     mQinvBins, mQinvVal[0], mQinvVal[1],
				     mPairKtBins, mPairKtVal[0], mPairKtVal[1]);

  //
  // Add correlation functions to the HbtManager
  //
  mPionMultAnalysis->AddCorrFctn(mPionQinvKtCF); 
  mPionMultAnalysis->AddCorrFctn(mPionBPLCMSFrameKtCF);   
  mPionMultAnalysis->SetNumEventsToMix(mEvents2Mix);   // set number of events to mix
  mHbtManager->AddAnalysis(mPionMultAnalysis); // add analysis to the manager
  
  std::cout << "---=== END CREATE ANALYSIS ===---\n";
  
  mChain->Init(); // StChain initialization

  //
  // Main loop over all events
  //
  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

  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 << "expPionHBT_pau200_y2015_3D[ERROR] - iRet: " << iRet << std::endl;
      break;
    }
  }

  mChain->Finish();
  
  TFile *oFile = new TFile(oFileName, "RECREATE"); // create output file
  std::cout << "Writing histograms to the output file...";

  
#ifdef MONITORS
  //
  // Event cut monitors
  //
  mFxtPionEventCutMonitorPass->VertexZ()->Write();
  mFxtPionEventCutMonitorPass->VertexYvsVertexX()->Write();
  mFxtPionEventCutMonitorPass->RefMult()->Write();
  mFxtPionEventCutMonitorPass->VpdVzDiff()->Write();
  mFxtPionEventCutMonitorFail->VertexZ()->Write();
  mFxtPionEventCutMonitorFail->VertexYvsVertexX()->Write();
  mFxtPionEventCutMonitorFail->RefMult()->Write();
  mFxtPionEventCutMonitorFail->VpdVzDiff()->Write();
  //
  // Positive pion track cut monitors
  //
  mFxtPionPosTrackCutMonitorPass->DCAGlobal()->Write();
  mFxtPionPosTrackCutMonitorPass->Nhits()->Write();
  mFxtPionPosTrackCutMonitorPass->P()->Write();
  mFxtPionPosTrackCutMonitorPass->Pt()->Write();
  mFxtPionPosTrackCutMonitorPass->PtVsNsigmaPion()->Write();
  mFxtPionPosTrackCutMonitorPass->PtVsNsigmaKaon()->Write();
  mFxtPionPosTrackCutMonitorPass->PtVsNsigmaProton()->Write();
  mFxtPionPosTrackCutMonitorPass->PvsDedx()->Write();
  mFxtPionPosTrackCutMonitorPass->Rapidity()->Write();
  mFxtPionPosTrackCutMonitorPass->PseudoRapidity()->Write();
  mFxtPionPosTrackCutMonitorPass->PtVsMassSqr()->Write();
  mFxtPionPosTrackCutMonitorFail->DCAGlobal()->Write();
  mFxtPionPosTrackCutMonitorFail->Nhits()->Write();
  mFxtPionPosTrackCutMonitorFail->P()->Write();
  mFxtPionPosTrackCutMonitorFail->Pt()->Write();
  mFxtPionPosTrackCutMonitorFail->PtVsNsigmaPion()->Write();
  mFxtPionPosTrackCutMonitorFail->PtVsNsigmaKaon()->Write();
  mFxtPionPosTrackCutMonitorFail->PtVsNsigmaProton()->Write();
  mFxtPionPosTrackCutMonitorFail->PvsDedx()->Write();
  mFxtPionPosTrackCutMonitorFail->Rapidity()->Write();
  mFxtPionPosTrackCutMonitorFail->PseudoRapidity()->Write();
  mFxtPionPosTrackCutMonitorFail->PtVsMassSqr()->Write();
  //
  // Negative pion track cut monitors
  //
  mFxtPionNegTrackCutMonitorPass->DCAGlobal()->Write();
  mFxtPionNegTrackCutMonitorPass->Nhits()->Write();
  mFxtPionNegTrackCutMonitorPass->P()->Write();
  mFxtPionNegTrackCutMonitorPass->Pt()->Write();
  mFxtPionNegTrackCutMonitorPass->PtVsNsigmaPion()->Write();
  mFxtPionNegTrackCutMonitorPass->PtVsNsigmaKaon()->Write();
  mFxtPionNegTrackCutMonitorPass->PtVsNsigmaProton()->Write();
  mFxtPionNegTrackCutMonitorPass->PvsDedx()->Write();
  mFxtPionNegTrackCutMonitorPass->Rapidity()->Write();
  mFxtPionNegTrackCutMonitorPass->PseudoRapidity()->Write();
  mFxtPionNegTrackCutMonitorPass->PtVsMassSqr()->Write();
  mFxtPionNegTrackCutMonitorFail->DCAGlobal()->Write();
  mFxtPionNegTrackCutMonitorFail->Nhits()->Write();
  mFxtPionNegTrackCutMonitorFail->P()->Write();
  mFxtPionNegTrackCutMonitorFail->Pt()->Write();
  mFxtPionNegTrackCutMonitorFail->PtVsNsigmaPion()->Write();
  mFxtPionNegTrackCutMonitorFail->PtVsNsigmaKaon()->Write();
  mFxtPionNegTrackCutMonitorFail->PtVsNsigmaProton()->Write();
  mFxtPionNegTrackCutMonitorFail->PvsDedx()->Write();
  mFxtPionNegTrackCutMonitorFail->Rapidity()->Write();
  mFxtPionNegTrackCutMonitorFail->PseudoRapidity()->Write();
  mFxtPionNegTrackCutMonitorFail->PtVsMassSqr()->Write();
  //
  // Pion pair cut monitors
  //
  mFxtPionPairCutMonitorPass->Kt()->Write();
  mFxtPionPairCutMonitorPass->Pt1Pt2Qinv()->Write();
  mFxtPionPairCutMonitorPass->FMRvsQinv()->Write();
  mFxtPionPairCutMonitorPass->SLvsQinv()->Write();
  mFxtPionPairCutMonitorPass->AveSepVsQinv()->Write();
  mFxtPionPairCutMonitorFail->Kt()->Write();
  mFxtPionPairCutMonitorFail->Pt1Pt2Qinv()->Write();
  mFxtPionPairCutMonitorFail->FMRvsQinv()->Write();
  mFxtPionPairCutMonitorFail->SLvsQinv()->Write();
  mFxtPionPairCutMonitorFail->AveSepVsQinv()->Write();
#endif

  //
  // Write correlation function to the
  //   output file
  //
  for (int iKt = 0; iKt < mPairKtBins; iKt++) {
    mPionQinvKtCF->Numerator(iKt)->Write();
    mPionQinvKtCF->Denominator(iKt)->Write();
      
    mPionBPLCMSFrameKtCF->Numerator(iKt)->Write();
    mPionBPLCMSFrameKtCF->Denominator(iKt)->Write();
    mPionBPLCMSFrameKtCF->QinvHisto(iKt)->Write();
  }
  std::cout << "\t[DONE]" << std::endl;

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

  //
  // Delete chain and prompt
  //   last message :D
  //
  delete mChain;

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