#include "StCosmicMuDstQAMaker.h"
#include "StBTofHeader.h"
#include "TBranch.h"

ClassImp(StCosmicMuDstQAMaker)

//
// Set maximum file size to 1.9 GB (Root has a 2GB limit)
//
#define MAXFILESIZE 1900000000



//_________________
StCosmicMuDstQAMaker::StCosmicMuDstQAMaker(StMuDstMaker *muMaker,
					   const Char_t *oFileName) {

  std::cout << "StCosmicMuDstQAMaker::StCosmicMuDstQAMaker - Creating an instance...";
  mMuDstMaker = muMaker;

  mEventIsGood = false;
  mNEventsIn = 0;
  mNEventsPassed = 0;

  mIsGoodTrack = false;
  mCurrentTrigger = 0;

  mFileName = oFileName;
  //
  // Initialize event cut variables
  //
  mPrimVtxZ[0] = -70.;
  mPrimVtxZ[1] = 70.;
  mPrimVtxR[0] = -1.;
  mPrimVtxR[1] = 10.;
  mPrimVtxVpdVzDiff[0] = -10.;
  mPrimVtxVpdVzDiff[1]= 10.;
  mPrimVtxXShift = 0.;
  mPrimVtxYShift = 0.;
  //
  // Initialize single-particle cut variables
  //
  mTrackP[0] = 0.1;
  mTrackP[1] = 2.;
  mTrackDca[0] = 0.;
  mTrackDca[1] = 5.;
  mTrackDcaGlobal[0] = 0.;
  mTrackDcaGlobal[1] = 5.;
  mTrackNHits[0] = 15;
  mTrackNHits[1] = 50;
  mTrackNHitsFit[0] = 15;
  mTrackNHitsFit[1] = 50;

  mTrackEta[0] = -1.1;
  mTrackEta[1] = 1.1;
  mTrackFlag[0] = 0;
  mTrackFlag[1] = 1000;

  /*
  //
  // Initialize TPC and TOF PID
  //
  mPionPionNSigma[0] = -3.;   mPionPionNSigma[1] = 3.;
  mPionKaonNSigma[0] = 1.;    mPionKaonNSigma[1] = -1.;
  */
  mTTTPThreshold = 0.7;
  
  std::cout << "\t[DONE]" << std::endl;
}

//_________________
StCosmicMuDstQAMaker::~StCosmicMuDstQAMaker() {
  /* nothing to do */
}

//_________________
Int_t StCosmicMuDstQAMaker::Init() {

  std::cout << "StCosmicMuDstQAMaker::Init - Initializing the maker"
	    << std::endl;

  //
  // Create histograms and output file
  //
  mOutFile = new TFile(mFileName, "RECREATE");
  //
  // General event distributions
  //
  hR = new TH1F("hR", "R=#sqrt{v_{x}^{2} + v_{y}^{2}};R (cm);#",
		200, 0., 2.);
  hVx = new TH1F("hVx", ";v_{x} (cm);#", 200, -1.5, 1.5);
  hVy = new TH1F("hVy", ";v_{y} (cm);#", 200, -1.5, 1.5);
  hVz = new TH1F("hVz", ";v_{z} (cm);#", 100, -70., 70.);
  hVxVsVy = new TH2F("hVxVsVy", ";v_{x} (cm); v_{y} (cm)",
		     600, -1.5, 1.5,
		     600, -1.5, 1.5);
  hVxVsVz = new TH2F("hVxVsVz", ";v_{x} (cm); v_{z} (cm)",
		     600, -1.5, 1.5,
		     600, -70., 70.);
  hVyVsVz = new TH2F("hVyVsVz", ";v_{y} (cm); v_{z} (cm)",
		     600, -1.5, 1.5,
		     600, -70., 70.);
  hNPrimTr = new TH1S("hNPrimTr", ";N_{primary tracks};#", 1000, 0, 1000);
  hNGlobTr = new TH1S("hNGlobTr", ";N_{global tracks};#", 1000, 0, 1000);
  hTofRefMult = new TH1S("hTofRefMult", ";TofRefMult;#", 1000, 0, 1000);
  hTofRefMultVsRefMult = new TH2S("hTofRefMultVsRefMult", ";TofRefMult;RefMult",
				  1000, 0., 1000,
				  1000, 0., 1000);
  hVpd = new TH1F("hVpd", ";v_{z}^{vpd} (cm);#", 100, -70., 70.);
  hVpdVz = new TH1F("hVpdVz", ";v_{z}^{vpd}-v_{z}", 100, -10., 10.);
  hNPrimVtx = new TH1C("hNPrimVtx", ";N_{primary vertex};#", 10, 0, 10);
  //
  // General track distributions
  //
  hTheta = new TH1F("hTheta", ";#theta (rad);#", 400, -0.5, 2.);
  hPhi = new TH1F("hPhi", ";#phi (rad);#", 800, -4., 4.);
  hThetaPhi = new TH2F("hThetaPhi", ";#phi (rad);#theta (rad)",
		       800, -4., 4.,
		       400, -0.5, 2.);
  hEtaPhi = new TH2F("hEtaPhi", ";#phi (rad);#eta",
		     800, -4., 4.,
		     400, -2., 2.);

  hP = new TH1F("hP", ";p (GeV/c);#", 3000, 0., 30.);  
  hPt = new TH1F("hPt", ";p_{t} (GeV/c);#", 3000, 0., 30.);
  hPx = new TH1F("hPx", ";p_{x} (GeV/c);#", 6000, -30, 30.);
  hPy = new TH1F("hPy", ";p_{y} (GeV/c);#", 6000, -30., 30.);
  hPz = new TH1F("hPz", ";p_{z} (GeV/c);#", 6000, -30., 30.);
  hEta = new TH1F("hEta", ";#eta;#", 400, -2., 2.);
  hMassSqrVsPt = new TH2F("hMassSqrVsPt", ";p_{t}Q (GeV/c);m^{2} (GeV/c^{2})",
			  400, -2., 2.,
			  1000, -0.1, 1.5);
  hDedxVsPt = new TH2F("hDedxVsPt", ";p_{t}Q (GeV/c);dE/dx (keV/cm)",
		       400, -2., 2.,
		       1000, 0., 10.);
  hInvBetaExpVsPt = new TH2F("hInvBetaExpVsPt", ";p_{t}Q (GeV/c);1/beta_{exp}",
			     400, -2., 2.,
			     200, 0., 2.);
  hTOF = new TH1F("hTOF", "time of flight;t ns;#", 100, 0., 50.);
  //
  // Any charge
  //
  hRefMult = new TH1F("hRefMult", ";RefMult;#", 2000, 0., 2000.);
  
  std::cout << "StCosmicMuDstQAMaker::Init - Initialization has been finished"
	    << std::endl;
  return StMaker::Init();
}

//________________
void StCosmicMuDstQAMaker::Clear(Option_t *option) {
  StMaker::Clear();
}

//________________
Int_t StCosmicMuDstQAMaker::Make() {

  mNEventsIn++;

  mMuDst = NULL;
  mMuEvent = NULL;

  //MuDstMaker initialization
  mMuDstMaker = (StMuDstMaker*)GetMaker("MuDst");
  if(!mMuDstMaker) {
    LOG_ERROR << "StCosmicMuDstQAMaker::Make [ERROR] - Cannot find StMuDstMaker"
	      << std::endl;
    return kStOk;
  }

  //Obtaining MuDst
  mMuDst = (StMuDst*)GetInputDS("MuDst");
  if(!mMuDst) {
    gMessMgr->Warning() << "StCosmicMuDstQAMaker::Make [WARNING] - No MuDst has been found" 
			<< endm;
    return kStOk;
  }

  //Obtaining MuEvent
  mMuEvent = (StMuEvent*)mMuDst->event();
  
  //Multiplicity cannot be negative
  unsigned short refMult = mMuEvent->refMult();
  unsigned short refMultPos = 0;
  unsigned short refMultNeg = 0;
  int mNVertices = mMuDst->numberOfPrimaryVertices();
  int mNPrimTracks = mMuDst->numberOfPrimaryTracks();
  int mNGlobTracks = mMuDst->numberOfGlobalTracks();

  if(refMult < 0) { //|| mNVertices < 0 || mNPrimTracks <= 0 || mNGlobTracks <= 0) {
    //std::cout << "StCosmicMuDstQAMaker::Make [WARNING] - bad event" << std::endl;
    return kStOk;
  }
  
  hRefMult->Fill(refMult);

  //Some initializations of local variables
  mPrimVertex = NULL;
  StThreeVectorF mVertPosition;
  Float_t mVpdVz = 0.;
  Int_t   mPrimVertIndex = -999;
  Float_t mRanking = -999.;


  //Clean index vectors
  CleanVariables();

  //Vertex loop
  unsigned short tofRefMult = 0;
  for(Int_t iVert=0; iVert<mNVertices; iVert++) {

    hNPrimVtx->Fill(mNVertices);

    mPrimVertex = mMuDst->primaryVertex(iVert);
    Float_t mVtxX = mVertPosition.x() - mPrimVtxXShift;
    Float_t mVtxY = mVertPosition.y() - mPrimVtxYShift;
    Float_t mVtxZ = mVertPosition.z();
    Float_t vtxR = TMath::Sqrt(mVtxX*mVtxX + mVtxY*mVtxY);
    Float_t vpdDiff = mVtxZ - mVpdVz;

    mPrimVertIndex = iVert;
    mRanking = mPrimVertex->ranking();
    mVertPosition = mPrimVertex->position();
    mVpdVz = mMuDst->btofHeader()->vpdVz();

    hR->Fill(vtxR);
    hVx->Fill(mVtxX);
    hVy->Fill(mVtxY);
    hVz->Fill(mVtxZ);
    hVxVsVy->Fill(mVtxX, mVtxY);
    hVxVsVz->Fill(mVtxX, mVtxZ);
    hVyVsVz->Fill(mVtxY, mVtxZ);
    hVpd->Fill(mVpdVz);
    hVpdVz->Fill(vpdDiff);
    
    mEventIsGood = true;

    hNPrimTr->Fill(mNPrimTracks);
    hNGlobTr->Fill(mNGlobTracks);
    //
    //Loop over primary tracks
    //
    for(int iTrk = 0; iTrk < mNGlobTracks; iTrk++) {
      //      mPrimTrack = mMuDst->primaryTracks(iTrk);
      mGlobTrack = mMuDst->globalTracks(iTrk);//mPrimTrack->index2Global());
      short charge = mGlobTrack->charge() > 0 ? 1 : -1;
      float eta = mGlobTrack->eta();
      float pt  = mGlobTrack->pt();
      float p   = mGlobTrack->p().mag();
      float p2  = mGlobTrack->p().mag2();
      float px  = mGlobTrack->p().x();
      float py  = mGlobTrack->p().y();
      float pz  = mGlobTrack->p().z();
      float theta = 2*TMath::ATan(TMath::Exp(-eta));
      float phi = mGlobTrack->phi();
      
      charge > 0 ? refMultPos++ : refMultNeg++;

      hP->Fill(p);

      hTheta->Fill(theta);
      hPhi->Fill(phi);
      hThetaPhi->Fill(phi, theta);
      hEtaPhi->Fill(phi, eta);

      hPt->Fill(pt);
      hPx->Fill(px);
      hPy->Fill(py);
      hPz->Fill(pz);
      hEta->Fill(eta);

      float dedx = mGlobTrack->dEdx();
      bool  tofTrack = IsTofTrack(mGlobTrack);
      float betaExp;
      float massSqr;

      hDedxVsPt->Fill(pt*charge, dedx*1e6);
      
      if (tofTrack) {
	hTOF->Fill(mGlobTrack->btofPidTraits().timeOfFlight());
	
	tofRefMult++;
	betaExp = mGlobTrack->btofPidTraits().beta();
	massSqr = p2*(1./(betaExp*betaExp) - 1.);
	hMassSqrVsPt->Fill(pt*charge, massSqr);

	hInvBetaExpVsPt->Fill(pt*charge, 1./betaExp);
      }
    } //for(Int_t iTrk=0; iTrk<mNPrimTracks; iTrk++)
  } //for(Int_t iVert=0; iVert<mNVertices; iVert++)
  
  hTofRefMult->Fill(tofRefMult);
  
  return kStOk;
}

//_________________
Bool_t StCosmicMuDstQAMaker::AcceptTrigger(StMuEvent *muEvent) {

  Bool_t mIsGoodTrigger = false;
  if(mTriggerIdCollection.empty()) {
    mIsGoodTrigger = true;
  }
  else {
    for(UInt_t iTrg=0; iTrg<mTriggerIdCollection.size(); iTrg++) {
      if(muEvent->triggerIdCollection().nominal().isTrigger(mTriggerIdCollection.at(iTrg))) {
	mIsGoodTrigger = true;
	mCurrentTrigger = mTriggerIdCollection.at(iTrg);
	break;
      }
    } //for(UInt_t iTrg=0; iTrg<mTriggerIdCollection.size(); iTrg++)
  }

  return mIsGoodTrigger;
}

//_________________
Bool_t StCosmicMuDstQAMaker::AcceptPrimVtx(StThreeVectorF vtxPos, 
					   Float_t vpdVz) {
  Float_t mVtxX = vtxPos.x() - mPrimVtxXShift;
  Float_t mVtxY = vtxPos.y() - mPrimVtxYShift;
  Float_t mVtxZ = vtxPos.z();
  Float_t vtxR = TMath::Sqrt(mVtxX*mVtxX + mVtxY*mVtxY);
  Float_t vpdDiff = mVtxZ - vpdVz;

  hR->Fill(vtxR);
  hVx->Fill(mVtxX);
  hVy->Fill(mVtxY);
  hVz->Fill(mVtxZ);
  hVxVsVy->Fill(mVtxX, mVtxY);
  hVxVsVz->Fill(mVtxX, mVtxZ);
  hVyVsVz->Fill(mVtxY, mVtxZ);
  hVpd->Fill(vpdVz);
  hVpdVz->Fill(vpdDiff);
  
  return ( vtxPos.z() >= mPrimVtxZ[0] &&
	   vtxPos.z() <= mPrimVtxZ[1] &&
	   vtxR >= mPrimVtxR[0] &&
	   vtxR <= mPrimVtxR[1] &&
	   vpdDiff >= mPrimVtxVpdVzDiff[0] &&
	   vpdDiff <= mPrimVtxVpdVzDiff[1] );
}

//_________________
Bool_t StCosmicMuDstQAMaker::AcceptTrack(StMuTrack *trk, UShort_t vtxInd) {
  float          mom      = trk->momentum().mag();
  float          eta      = trk->eta();
  unsigned short nHitsFit = trk->nHitsFit();
  short          flag     = trk->flag();
  float          dca      = trk->dcaGlobal(vtxInd).perp();
  
  bool isMomOk      = ( mom >= mTrackP[0] &&
			mom <= mTrackP[1] );
  bool isEtaOk      = ( eta >= mTrackEta[0] &&
			eta <= mTrackEta[1] );
  bool isNHitsFitOk = ( nHitsFit >= mTrackNHitsFit[0] &&
			nHitsFit <= mTrackNHitsFit[1] );
  bool isFlagOk     = ( flag >= mTrackFlag[0] &&
			flag <= mTrackFlag[1] );
  bool isDcaOk      = ( dca >= mTrackDcaGlobal[0] &&
			dca <= mTrackDcaGlobal[1] );

  return ( isMomOk && isEtaOk && isNHitsFitOk && isFlagOk && isDcaOk );
}


//_________________
Int_t StCosmicMuDstQAMaker::Finish() {

  std::cout << "*************************************" << std::endl
	    << "StCosmicMuDstQAMaker has been finished" << std::endl
	    << "\t nEventsPassed   : " << mNEventsPassed  << std::endl
	    << "\t nEventsProcessed: " << mNEventsIn << std::endl
	    << "*************************************" << std::endl;
  mOutFile->cd();
  mOutFile->Write();
  mOutFile->Close();
  
  return kStOk;
}

//_________________
void StCosmicMuDstQAMaker::CleanVariables() {

  mEventIsGood = false;
}

//_________________
void StCosmicMuDstQAMaker::SetTriggerId(unsigned int id) {
  mTriggerIdCollection.push_back(id);
}


//_________________
void StCosmicMuDstQAMaker::SetMuDstMaker(StMuDstMaker *maker) {
  mMuDstMaker = maker;
}


//_________________
void StCosmicMuDstQAMaker::SetVtxZCut(float lo, float hi) {
  mPrimVtxZ[0] = lo;
  mPrimVtxZ[1] = hi;
}


//_________________
void StCosmicMuDstQAMaker::SetVtxRCut(float lo, float hi) {
  mPrimVtxR[0] = lo;
  mPrimVtxR[1] = hi;  
}


//_________________
void StCosmicMuDstQAMaker::SetVtxShift(float xShift, float yShift) {
  mPrimVtxXShift = xShift;
  mPrimVtxYShift = yShift;
}


//_________________
void StCosmicMuDstQAMaker::SetVtxVpdVzDiffCut(float lo, float hi) {
  mPrimVtxVpdVzDiff[0] = lo;
  mPrimVtxVpdVzDiff[1] = hi;
}


//_________________
void StCosmicMuDstQAMaker::SetP(float lo, float hi) {
  mTrackP[0] = lo;
  mTrackP[1] = hi;
}


//////////////////////////////////////////////////////////
//_________________
void StCosmicMuDstQAMaker::SetPhi(float phi) {
  mPhi = phi;
}

//______________
void StCosmicMuDstQAMaker::SetTheta(float theta) {
  mTheta = theta;
}
//////////////////////////////////////////////////////////
//_________________
void StCosmicMuDstQAMaker::SetPt(float lo, float hi) {
  mTrackPt[0] = lo;
  mTrackPt[1] = hi;
}


//_________________
void StCosmicMuDstQAMaker::SetTrackNHits(int lo, int hi) {
  mTrackNHits[0] = lo;
  mTrackNHits[1] = hi;  
}


//_________________
void StCosmicMuDstQAMaker::SetTrackNHitsFit(int lo, int hi) {
  mTrackNHitsFit[0] = lo;
  mTrackNHitsFit[1] = hi;
}


//_________________
void StCosmicMuDstQAMaker::SetTrackEta(float lo, float hi) {
  mTrackEta[0] = lo;
  mTrackEta[1] = hi;
}


//_________________
void StCosmicMuDstQAMaker::SetTrackFlag(short lo, short hi) {
  mTrackFlag[0] = lo;
  mTrackFlag[1] = hi;
}


//_________________
void StCosmicMuDstQAMaker::SetTrackDCA(float lo, float hi) {
  mTrackDcaGlobal[0] = lo;
  mTrackDcaGlobal[1] = hi;
}


//_________________
Bool_t StCosmicMuDstQAMaker::IsTofTrack(StMuTrack *trk) { //Only for globals
  Bool_t isTofHit = false;
  if(trk->btofPidTraits().beta() > 0 &&
     trk->btofPidTraits().timeOfFlight() > 0) {
    isTofHit = true;
  }
  return isTofHit;
}


//_________________
void StCosmicMuDstQAMaker::SetPionPionNSigma(float lo, float hi) {
  mPionPionNSigma[0] = lo;   mPionPionNSigma[1] = hi;
}


//_________________
void StCosmicMuDstQAMaker::SetPionKaonNSigma(float lo, float hi) {
  mPionKaonNSigma[0] = lo;   mPionKaonNSigma[1] = hi;
}


//_________________
void StCosmicMuDstQAMaker::SetPionProtonNSigma(float lo, float hi) {
  mPionProtonNSigma[0] = lo;   mPionProtonNSigma[1] = hi;
}


//_________________
void StCosmicMuDstQAMaker::SetKaonPionNSigma(float lo, float hi) {
  mKaonPionNSigma[0] = lo;   mKaonPionNSigma[1] = hi;
}


//_________________
void StCosmicMuDstQAMaker::SetKaonKaonNSigma(float lo, float hi) {
  mKaonKaonNSigma[0] = lo;   mKaonKaonNSigma[1] = hi;
}


//_________________
void StCosmicMuDstQAMaker::SetKaonProtonNSigma(float lo, float hi) {
  mKaonProtonNSigma[0] = lo;   mKaonProtonNSigma[1] = hi;
}


//_________________
void StCosmicMuDstQAMaker::SetProtonPionNSigma(float lo, float hi) {
  mProtonPionNSigma[0] = lo;   mProtonPionNSigma[1] = hi;
}


//_________________
void StCosmicMuDstQAMaker::SetProtonKaonNSigma(float lo, float hi) {
  mProtonKaonNSigma[0] = lo;   mProtonKaonNSigma[1] = hi;
}


//_________________
void StCosmicMuDstQAMaker::SetProtonProtonNSigma(float lo, float hi) {
  mProtonProtonNSigma[0] = lo;   mProtonProtonNSigma[1] = hi;
}


//_________________
void StCosmicMuDstQAMaker::SetPionMassSqr(float lo, float hi) {
  mPionMass[0] = lo;   mPionMass[1] = hi;
}


//_________________
void StCosmicMuDstQAMaker::SetKaonMassSqr(float lo, float hi) {
  mKaonMass[0] = lo;   mKaonMass[1] = hi;
}


//_________________
void StCosmicMuDstQAMaker::SetProtonMassSqr(float lo, float hi) {
  mProtonMass[0] = lo;   mProtonMass[1] = hi;
}


//_________________
void StCosmicMuDstQAMaker::SetTTTPThreshold(float pThresh) {
  mTTTPThreshold = pThresh;
}