//
//  The version of the class should be changed every time
//  when any changes in the codes are done
//  Grigory Nigmatkulov: 2016/12/15
//
#ifndef STFEMTODSTINCLUSIVESELECTOR_HH
#define STFEMTODSTINCLUSIVESELECTOR_HH

#include "StMaker.h"
#include "StThreeVector.hh"
#include "StThreeVectorF.hh"
#include "StPhysicalHelixD.hh"
#include "StRefMultCorr/StRefMultCorr.h"

#include "StFemtoEvent.h"
#include "StFemtoTrack.h"

#include "StMuDSTMaker/COMMON/StMuDstMaker.h"
#include "StMuDSTMaker/COMMON/StMuDst.h"
#include "StMuDSTMaker/COMMON/StMuEvent.h"
#include "StMuDSTMaker/COMMON/StMuTrack.h"
#include "StMuDSTMaker/COMMON/StMuPrimaryVertex.h"
#include "StMuDSTMaker/COMMON/StMuBTofPidTraits.h"
#include "StMuDSTMaker/COMMON/StMuBTofHit.h"

#include <TString.h>
#include <TFile.h>
#include <TLorentzVector.h>
#include <TTree.h>
#include <TH1.h>
#include <TH2.h>

#include <vector>
#include <iostream>

#ifndef St_NO_NAMESPACES
using namespace std;
#endif

//Predefinisions
class StFemtoEvent;
class StFemtoTrack;

class StMuDst;
class StMuDstMaker;
class StMuEvent;
class StMuTrack;

//_________________
class StFemtoDstInclusiveSelector : public StMaker {
 public:
  StFemtoDstInclusiveSelector(StMuDstMaker *muMaker,
			      const Char_t *oFileName);
  ~StFemtoDstInclusiveSelector();

  void Clear(Option_t *option = "");
  Int_t Init();
  Int_t Make();
  Int_t Finish();

  void CleanVariables();

  void SetCompressionLevel(Int_t comp)        { mCompression = comp; }

  //Setters
  //Global
  void SetCollisionTypeAuAu(const Bool_t& type)   { mCollisionTypeAuAu = type; }
  void SetMuDstMaker(StMuDstMaker *maker)     { mMuDstMaker = maker; }
  void SetRefmultCorrUtil(StRefMultCorr* rf)  { mRefMultCorrUtil = rf; }
  void SetAuAuZdcCoincidenceEnergy(const Bool_t &correct) { mAuAuCorrectZdc = correct; }

  //Event
  void AddTriggerId(const UInt_t& id);
  void SetVtxZCut(const Float_t& lo, const Float_t& hi)   { mPrimVtxZ[0]=lo; mPrimVtxZ[1]=hi; }
  void SetVtxRCut(const Float_t& lo, const Float_t& hi)   { mPrimVtxR[0]=lo; mPrimVtxR[1]=hi; }
  void SetVtxShift(const Float_t& xShift, const Float_t& yShift) { mPrimVtxXShift=xShift; mPrimVtxYShift=yShift; }
  void SetVtxVpdVzDiffCut(const Float_t& lo, const Float_t& hi)  { mPrimVtxVpdVzDiff[0]=lo; mPrimVtxVpdVzDiff[1]=hi;}

  //Track
  void SetParticleMomentum(const Float_t& lo, const Float_t& hi) { mTrackMomentum[0]=lo; mTrackMomentum[1]=hi; }
  void SetTrackDca(const Float_t& lo, const Float_t& hi)  { mTrackDca[0]=lo; mTrackDca[1]=hi; }
  void SetTrackDcaGlobal(const Float_t& lo, const Float_t& hi) { mTrackDcaGlobal[0]=lo; mTrackDcaGlobal[1]=hi; }
  void SetTrackNHits(const Int_t& lo, const Int_t& hi)    { mTrackNHits[0]=lo; mTrackNHits[1]=hi; }
  void SetTrackEta(const Float_t& lo, const Float_t& hi)  { mTrackEta[0]=lo; mTrackEta[1]=hi; }
  void SetTrackFlag(const Short_t& lo, const Short_t& hi) { mTrackFlag[0]=lo; mTrackFlag[1]=hi;}

  //Inclusive cuts
  //Pions
  void SetSelectPions(Bool_t &sel) { mSelectPions = sel; }
  void SetSelectTpcPions(Bool_t &sel) { mSelectTpcPions = sel; }
  void SetPionTpcMom(Float_t &lo, Float_t &hi) { mPionTpcMomentum[0]=lo; mPionTpcMomentum[1]=hi; }
  void SetPionTpcNSigmaElectron(Float_t &lo, Float_t &hi) { mPionTpcNSigmaElectron[0]=lo; mPionTpcNSigmaElectron[1]=hi; }
  void SetPionTpcNSigmaPion(Float_t &lo, Float_t &hi) { mPionTpcNSigmaPion[0]=lo; mPionTpcNSigmaPion[1]=hi; }
  void SetPionTpcNSigmaKaon(Float_t &lo, Float_t &hi) { mPionTpcNSigmaKaon[0]=lo; mPionTpcNSigmaKaon[1]=hi; }
  void SetPionTpcNSigmaProton(Float_t &lo, Float_t &hi) { mPionTpcNSigmaProton[0]=lo; mPionTpcNSigmaProton[1]=hi; }
  void SetSelectTofPions(Bool_t &sel) {mSelectTofPions = sel; }
  void SetPionTofMom(Float_t &lo, Float_t &hi) { mPionTofMomentum[0]=lo; mPionTofMomentum[1]=hi; }
  void SetPionMassSqr(Float_t &lo, Float_t &hi) { mPionMassSqr[0]=lo; mPionMassSqr[1]=hi; }

  //Kaons
  void SetSelectKaons(Bool_t &sel) { mSelectKaons = sel; }
  void SetSelectTpcKaons(Bool_t &sel) { mSelectTpcKaons = sel; }
  void SetKaonTpcMom(Float_t &lo, Float_t &hi) { mKaonTpcMomentum[0]=lo; mKaonTpcMomentum[1]=hi; }
  void SetKaonTpcNSigmaElectron(Float_t &lo, Float_t &hi) { mKaonTpcNSigmaElectron[0]=lo; mKaonTpcNSigmaElectron[1]=hi; }
  void SetKaonTpcNSigmaPion(Float_t &lo, Float_t &hi) { mKaonTpcNSigmaPion[0]=lo; mKaonTpcNSigmaPion[1]=hi; }
  void SetKaonTpcNSigmaKaon(Float_t &lo, Float_t &hi) { mKaonTpcNSigmaKaon[0]=lo; mKaonTpcNSigmaKaon[1]=hi; }
  void SetKaonTpcNSigmaProton(Float_t &lo, Float_t &hi) { mKaonTpcNSigmaProton[0]=lo; mKaonTpcNSigmaProton[1]=hi; }
  void SetSelectTofKaons(Bool_t &sel) {mSelectTofKaons = sel; }
  void SetKaonTofMom(Float_t &lo, Float_t &hi) { mKaonTofMomentum[0]=lo; mKaonTofMomentum[1]=hi; }
  void SetKaonMassSqr(Float_t &lo, Float_t &hi) { mKaonMassSqr[0]=lo; mKaonMassSqr[1]=hi; }

  //Protons
  void SetSelectProtons(Bool_t &sel) { mSelectProtons = sel; }
  void SetSelectTpcProtons(Bool_t &sel) { mSelectTpcProtons = sel; }
  void SetProtonTpcMom(Float_t &lo, Float_t &hi) { mProtonTpcMomentum[0]=lo; mProtonTpcMomentum[1]=hi; }
  void SetProtonTpcNSigmaElectron(Float_t &lo, Float_t &hi) { mProtonTpcNSigmaElectron[0]=lo; mProtonTpcNSigmaElectron[1]=hi; }
  void SetProtonTpcNSigmaPion(Float_t &lo, Float_t &hi) { mProtonTpcNSigmaPion[0]=lo; mProtonTpcNSigmaPion[1]=hi; }
  void SetProtonTpcNSigmaKaon(Float_t &lo, Float_t &hi) { mProtonTpcNSigmaKaon[0]=lo; mProtonTpcNSigmaKaon[1]=hi; }
  void SetProtonTpcNSigmaProton(Float_t &lo, Float_t &hi) { mProtonTpcNSigmaProton[0]=lo; mProtonTpcNSigmaProton[1]=hi; }
  void SetSelectTofProtons(Bool_t &sel) {mSelectTofProtons = sel; }
  void SetProtonTofMom(Float_t &lo, Float_t &hi) { mProtonTofMomentum[0]=lo; mProtonTofMomentum[1]=hi; }
  void SetProtonMassSqr(Float_t &lo, Float_t &hi) { mProtonMassSqr[0]=lo; mProtonMassSqr[1]=hi; }

 private:

  const Char_t *mOutFileName;
  TFile *mOutFile;
  TTree *mTree;
  StMuDstMaker *mMuDstMaker;
  StMuDst *mMuDst;
  Int_t mCompression;
  StRefMultCorr *mRefMultCorrUtil;
  
  Int_t  mNBytes;
  vector<UInt_t> mTriggerIdCollection;

  //Global counters
  Bool_t mEventIsGood;
  Int_t  mNEventsIn;
  Int_t  mNEventsPassed;

  Bool_t mIsGoodTrack;
  Bool_t mAuAuCorrectZdc;
  Bool_t mLoopAllVert;

  //Event cuts
  Bool_t  mCollisionTypeAuAu;        // pp: false; AuAu: true
  Float_t mPrimVtxZ[2];
  Float_t mPrimVtxR[2];
  Float_t mPrimVtxVpdVzDiff[2];
  Float_t mPrimVtxXShift;
  Float_t mPrimVtxYShift;

  //Single-particle cuts
  Float_t  mTrackMomentum[2];
  Float_t  mTrackDca[2];
  Float_t  mTrackDcaGlobal[2];
  UShort_t mTrackNHits[2];
  Float_t  mTrackEta[2];
  Short_t  mTrackFlag[2];

  //Inclusive cuts
  //Pions
  Bool_t mSelectPions;
  Bool_t mSelectTpcPions;
  Float_t mPionTpcMomentum[2];
  Float_t mPionTpcNSigmaElectron[2];
  Float_t mPionTpcNSigmaPion[2];
  Float_t mPionTpcNSigmaKaon[2];
  Float_t mPionTpcNSigmaProton[2];
  Bool_t mSelectTofPions;
  Float_t mPionMassSqr[2];
  Float_t mPionTofMomentum[2];

  //Kaons
  Bool_t mSelectKaons;
  Bool_t mSelectTpcKaons;
  Float_t mKaonTpcMomentum[2];
  Float_t mKaonTpcNSigmaElectron[2];
  Float_t mKaonTpcNSigmaPion[2];
  Float_t mKaonTpcNSigmaKaon[2];
  Float_t mKaonTpcNSigmaProton[2];
  Bool_t mSelectTofKaons;
  Float_t mKaonTofMomentum[2];
  Float_t mKaonMassSqr[2];

  //Protons
  Bool_t mSelectProtons;
  Bool_t mSelectTpcProtons;
  Float_t mProtonTpcMomentum[2];
  Float_t mProtonTpcNSigmaElectron[2];
  Float_t mProtonTpcNSigmaPion[2];
  Float_t mProtonTpcNSigmaKaon[2];
  Float_t mProtonTpcNSigmaProton[2];
  Bool_t mSelectTofProtons;
  Float_t mProtonTofMomentum[2];
  Float_t mProtonMassSqr[2];

  //Trigger is a bit mask for the list of triggers that is set in macros
  UInt_t       mCurrentTrigger;

  // Femto event structure
  StFemtoEvent *mFemtoEvent;
  // Femto track
  StFemtoTrack *mFTrack;
  
  //Internal class methods
  Bool_t AcceptTrigger(StMuEvent *muEvent);
  Bool_t AcceptRefMult2(StMuTrack *trk);
  Bool_t AcceptPrimVtx(StThreeVectorF vtxPos, Float_t vpdVz);
  Bool_t AcceptRunId(Int_t runId);   //For StRefMultCorr (isBadRun)
  Bool_t AcceptTrack(StMuTrack *primTrk, UShort_t vtxInd);
  Bool_t AcceptDcaGlobal(Float_t dcaToPrimVtx);
  Bool_t IsTofTrack(StMuTrack *globTrack);

  //pion, kaon and proton selection
  Bool_t IsTpcPion(StMuTrack *pTrk);
  Bool_t IsTofPion(Float_t mSqr, Float_t p);
  Bool_t IsTpcKaon(StMuTrack *pTrk);
  Bool_t IsTofKaon(Float_t mSqr, Float_t p);
  Bool_t IsTpcProton(StMuTrack *pTrk);
  Bool_t IsTofProton(Float_t mSqr, Float_t p);


  ClassDef(StFemtoDstInclusiveSelector, 3)
};

#endif