#include "StHbtO97DstReader.h"

ClassImp(StHbtO97DstReader)

const Float_t PYTHIA_RBFIELD = -4.5106000;

//_________________
StHbtO97DstReader::StHbtO97DstReader(const char *dir, const char *fileName,
				     const char *filter, int maxFiles, bool debug) {
  mDir      = string(dir);
  mFileName = string(fileName);
  mFilter   = filter;
  mMaxFiles = maxFiles;
  mDebug    = debug;

  mTotalTracks = 0.;

  mTree       = 0;
  mTChain     = 0;
  mO97Event   = 0;
  mEventIndex = 0;

  InitRead(mDir, mFileName, mFilter, mMaxFiles);
  mNEvents = (unsigned int)mTChain->GetEntries();
}

//_________________
int StHbtO97DstReader::InitRead(string dir, string fileName,
				string filter, int maxFiles) {
  mEventIndex = 0;
  mTChain = new TChain("StO97Dst", "StO97Dst");

  std::cout << "[StHbtO97DstReader]: ---=== StHbtO97DstReader initialization ===---" << std::endl;

  int numFiles = 0;
  if (!fileName.empty()) {

    if (strstr(fileName.c_str(), ".lis") ||
	strstr(fileName.c_str(), ".list")) {
      numFiles = FillChain(mTChain, 
			   (dir + fileName).c_str(), 
			   maxFiles);
    }
    else {
      mTChain->Add((dir + fileName).c_str());
      numFiles++;
    }
  } 
  else {
    numFiles = FillChain(mTChain, (char *)dir.c_str(), // not cool
			 filter.c_str(), maxFiles);
  }
  mTChain->SetBranchAddress("StO97Event", &mO97Event);
  mNEvents = (unsigned int)mTChain->GetEntries();
  return numFiles;
}

//_________________
int StHbtO97DstReader::FillChain(TChain *chain, const char *fileName, 
				 int maxFiles) {

  std::ifstream inStream(fileName);
  if (!inStream.is_open()) {
    std::cout << "[StHbtO97DstReader]: can't open file list: "
	      << fileName << std::endl;
    return 0;
  }
  std::cout << "[StHbtO97DstReader]: using file list: "
	    << fileName << std::endl;

  int count = 0;
  char buf[0xFF];
  for ( ; inStream.good(); ) {
    inStream.getline(buf, 0xFF);
    TString fileName(buf);
    if (fileName.Contains("root")) {
      chain->Add(buf);
      count++;
      if ((maxFiles > 0) && (count > maxFiles)) break;
    }
  }

  std::cout << "[StHbtO97DstReader]: Added " << count
	    << " files to the chain" << std::endl;
  return count;
}

//_________________
int StHbtO97DstReader::FillChain(TChain *chain, char *dir,
				 const char *filter, int maxFiles) {

  TSystem *gSystem = new TSystem();
  if (!gSystem) {
    std::cout << "[StHbtO97DstReader]: cannot allocate memory for TSystem"
	      << std::endl;
    return 0;
  }
  void *pdir = gSystem->OpenDirectory(dir);
  if (!pdir) {
    std::cout << "[StHbtO97DstReader]: can't open directory " 
	      << dir << std::endl;
    delete gSystem;
    return 0;
  }
  
  int count = 0;
  char *fileName;

  while ((fileName = (char *)gSystem->GetDirEntry(dir))) {
    if (strcmp(fileName, ".") == 0 || strcmp(fileName, "..") == 0)
      continue;
    if (strstr(fileName, filter)) {
      char *fullName = gSystem->ConcatFileName(dir, fileName);
      std::cout << "[StHbtO97DstReader]: Adding "
		<< fullName << " to the chain" << std::endl;
      chain->Add(fullName);
      delete fullName;
      count++;
      if ((maxFiles > 0) && (count > maxFiles)) break;
    }
  }
  
  std::cout << "[StHbtO97DstReader]: Added " << count 
	    << " files to the chain" << std::endl;
  delete gSystem;
  return count;
}

//_________________
void StHbtO97DstReader::UninitRead() {
  if (mTChain)     delete mTChain;
  if (mO97Event)   delete mO97Event;
  
  mO97Event = 0;
  mTChain   = 0;
}

//_________________
float StHbtO97DstReader::dedxMean(float mass, float mom) {
  float dedxMean;
  float tpcDedxGain   = 0.174325e-06;
  float tpcDedxOffset = -2.71889; 
  float tpcDedxRise   = 776.626;
  
  float gamma = TMath::Sqrt(mom*mom/(mass*mass) + 1.);
  float beta  = TMath::Sqrt(1. - 1./(gamma*gamma));
  float rise  = tpcDedxRise*beta*beta*gamma*gamma;
  if (beta > 0)
    dedxMean = tpcDedxGain/(beta*beta)*(0.5*TMath::Log(rise) - beta*beta 
					- tpcDedxOffset);
  else
    dedxMean = 1000.;
  return dedxMean;
}

//_________________
StHbtEvent *StHbtO97DstReader::Read() {

  if (mDebug) {
    std::cout << "[StHbtO97DstReader]: cEventIndex = "
	      << mEventIndex << " cNEvents = "
	      << mNEvents << std::endl;
  }
      
  if (!mNEvents) {
    std::cout << "[StHbtO97DstReader]: no events to read\n";
    return 0;
  }

  mO97Event->Clear();
  int bytes = mTChain->GetEntry(mEventIndex++);

  if (mNEvents < mEventIndex) {
    std::cout << "[StHbtO97DstReader]: EOF\n";
    return 0;
  }

  if (!bytes) {
    std::cout << "[StHbtO97DstReader]: no event\n";
    return 0;
  }

  StHbtEvent *mHbtEvent = 0;

  if (mO97Event) {
    mHbtEvent = new StHbtEvent();

    int nTracks = mO97Event->GetNTracks();
    int nPosTracks = 0;
    int nNegTracks = 0;
    StThreeVectorF vPos( 0., 0., 0. );
    
    TClonesArray *mO97TrackArr = mO97Event->GetTracks();
    
    if(mO97TrackArr) {
      mTotalTracks += (float)nTracks;
      if (mDebug) {
	std::cout << "[StHbtO97DstReader]: nTracks = "
		  << nTracks << std::endl;
      }
      
      for (int i = 0; i < nTracks; i++) {

	// Initialize StO97Track and StHbtTrack
	StO97Track *mO97Track = (StO97Track *)mO97TrackArr->At(i);
	StHbtTrack *mHbtTrack = new StHbtTrack();

	// Set charge, mass and momentum
	char charge = mO97Track->GetPdgId()/abs(mO97Track->GetPdgId());
	if (charge > 0) {
	  nPosTracks++;
	}
	else {
	  nNegTracks++;
	}

	float trkMass = mO97Track->GetMass();
	StThreeVectorF trkMom(mO97Track->GetPx(), 
			      mO97Track->GetPy(), 
			      mO97Track->GetPz());

	// Set nSigma and PID probability
	float trkNSigmaElectron,  trkNSigmaPion;
	float trkNSigmaKaon,      trkNSigmaProton;
	float trkPidProbElectron, trkPidProbPion;
	float trkPidProbKaon,     trkPidProbProton;	
	switch (abs(mO97Track->GetPdgId())) {
	case 211: // pion
	  trkNSigmaElectron = -99.;
	  trkNSigmaPion     = 0.;
	  trkNSigmaKaon     = -99.;
	  trkNSigmaProton   = -99.;
	  trkPidProbElectron = -99.;
	  trkPidProbPion     = 0.;
	  trkPidProbKaon     = -99.;
	  trkPidProbProton   = -99.;
	  break;
	case 321: // kaon
	  trkNSigmaElectron = -99.;
	  trkNSigmaPion     = -99.;
	  trkNSigmaKaon     = 0.;
	  trkNSigmaProton   = -99.;
	  trkPidProbElectron = -99.;
	  trkPidProbPion     = -99.;
	  trkPidProbKaon     = 0.;
	  trkPidProbProton   = -99.;
	  break;
	case 2212: // proton
	  trkNSigmaElectron = -99.;
	  trkNSigmaPion     = -99.;
	  trkNSigmaKaon     = -99.;
	  trkNSigmaProton   = 0.;
	  trkPidProbElectron = -99.;
	  trkPidProbPion     = -99.;
	  trkPidProbKaon     = -99.;
	  trkPidProbProton   = 0.;
	  break;
	default:
	  trkNSigmaElectron = -99.;
	  trkNSigmaPion     = -99.;
	  trkNSigmaKaon     = -99.;
	  trkNSigmaProton   = -99.;
	  trkPidProbElectron = -99.;
	  trkPidProbPion     = -99.;
	  trkPidProbKaon     = -99.;
	  trkPidProbProton   = -99.;
	}

	// Set mO97Track info
	mHbtTrack->SetTrackId( mO97Track->GetId() );
	//mHbtTrack->SetTrackType( 1 ); // 0 - global, 1 - primary
	//mHbtTrack->SetCharge( charge );
	mHbtTrack->SetFlag( 400 );
	
	mHbtTrack->SetNHits( charge * 45 );
	//mHbtTrack->SetNHitsFit( 45 );
	mHbtTrack->SetNHitsPossible( 45 );
	mHbtTrack->SetNHitsDedx( 45 );
	
	mHbtTrack->SetNSigmaElectron( trkNSigmaElectron );
	mHbtTrack->SetNSigmaPion( trkNSigmaPion );
	mHbtTrack->SetNSigmaKaon( trkNSigmaKaon );
	mHbtTrack->SetNSigmaProton( trkNSigmaProton );
	
	mHbtTrack->SetPidProbElectron( trkPidProbElectron );
	mHbtTrack->SetPidProbPion( trkPidProbPion );
	mHbtTrack->SetPidProbKaon( trkPidProbKaon );
	mHbtTrack->SetPidProbProton( trkPidProbProton );
	
	mHbtTrack->SetdEdx( dedxMean(trkMass, trkMom.mag()) );

	// Set DCA in cm
	StThreeVectorF origin( mO97Track->GetXfr()*1e-13,
			       mO97Track->GetYfr()*1e-13,
			       mO97Track->GetZfr()*1e-13 );
	float dcaXY = origin.perp();
	float dcaZ = origin.z();
	mHbtTrack->SetDCAxyGlobal( dcaXY );
	mHbtTrack->SetDCAzGlobal( dcaZ );
	mHbtTrack->SetChi2(1.);

	if( (dcaXY*dcaXY + dcaZ*dcaZ) > 2. ) { //Lets assume that if dca>2cm -> global track
	  mHbtTrack->SetP( StThreeVectorF(0., 0., 0.) );
	}
	else {
	  mHbtTrack->SetP( trkMom );
	}
	mHbtTrack->SetPGlobal( trkMom );
	StPhysicalHelixD helixGlob( trkMom, origin, PYTHIA_RBFIELD * kilogauss, charge );
	StPhysicalHelixD helixPrim( trkMom, origin, PYTHIA_RBFIELD * kilogauss, charge );
	mHbtTrack->SetHelixGlobal( helixGlob );
	mHbtTrack->SetHelix( helixPrim );
	
	mHbtTrack->SetTopologyMap( 0, 0. );
	mHbtTrack->SetTopologyMap( 1, 0. );
	
	mHbtTrack->SetTofBeta( trkMom.mag()/mO97Track->GetEnergy() );
	mHbtTrack->SetHiddenInfo(0);
	/*
	mHbtTrack->SetFreezeOutX( mO97Track->GetXfr() );
	mHbtTrack->SetFreezeOutY( mO97Track->GetYfr() );
	mHbtTrack->SetFreezeOutZ( mO97Track->GetZfr() );
	mHbtTrack->SetFreezeOutT( mO97Track->GetTfr() );
	*/
	
	mHbtEvent->TrackCollection()->push_back( mHbtTrack );
      } // for (int i = 0; i < nTracks; i++)
    } // if (mO97TrackArr)

    mHbtEvent->SetEventNumber( mO97Event->GetEventNumber() );
    mHbtEvent->SetRunNumber( 0 );
    mHbtEvent->SetZdcAdcEast( 45000 );
    mHbtEvent->SetZdcAdcWest( 45000 );
    mHbtEvent->SetCtbMult( 0 );
    mHbtEvent->SetNumberOfTpcHits( 5000 );
    mHbtEvent->SetNumberOfGoodTracks( nTracks );
    mHbtEvent->SetReactionPlane( 0., 0 );
    mHbtEvent->SetReactionPlaneError( 0., 0 );
    mHbtEvent->SetReactionPlaneSubEventDifference( 0., 0 );
    mHbtEvent->SetUncorrectedNumberOfPositivePrimaries( nPosTracks );
    mHbtEvent->SetUncorrectedNumberOfNegativePrimaries( nNegTracks );
    mHbtEvent->SetUncorrectedNumberOfPrimaries( nTracks );
    mHbtEvent->SetPrimVertPos( vPos );
    mHbtEvent->SetMagneticField( PYTHIA_RBFIELD );
    mHbtEvent->SetVpdVz( 0. );
    mHbtEvent->SetSimulation(true);
  } // if (mO97Event)

  return mHbtEvent;
}

//_________________
StHbtEvent *StHbtO97DstReader::ReturnHbtEvent() {

  StHbtEvent *mHbtEvent = Read();
  if (!mHbtEvent) {
    std::cout << "[StHbtO97DstReader]: no hbt event (mHbtEvent == NULL)"
	      << std::endl;
    return 0;
  }
  else {
    return mHbtEvent;
  }
}

//_________________
StHbtString StHbtO97DstReader::Report() {
  StHbtString temp = "\nThis is the StHbtO97DstReader\n";
  return temp;
}

//_________________
void StHbtO97DstReader::PrintTotalTracks() {
  std::cout << "[StHbtO97DstReader]: TotalTracks = "
	    << mTotalTracks << std::endl;
}

//_________________
int StHbtO97DstReader::GetNEvents() {
  if (mTChain) {
    return mNEvents;
  }
  else {
    return -1;
  }
}
