#include "StHbtFemtoDstReader.h"
#include "TString.h"

ClassImp(StHbtFemtoDstReader)

//_________________
StHbtFemtoDstReader::StHbtFemtoDstReader(const char *aDir, 
					 const char *aFileName,
					 const char *aFilter,
					 int aMaxFiles,
					 bool aDebug) {
  mDir      = string(aDir);
  mFileName = string(aFileName);
  mFilter   = aFilter;
  mMaxFiles = aMaxFiles;
  mDebug    = aDebug;

  mTotalTracks = 0.;
  
  mTree       = 0;
  mTChain     = 0;
  mFemtoEvent = 0;
  mEventIndex = 0;
  mTrig       = 0;

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

//_________________
int StHbtFemtoDstReader::FillChain(TChain *aChain, char *aDir,
				   const char *aFilter, int aMaxFiles) {

  TSystem *gSystem = new TSystem();
  if (!gSystem) {
    std::cout << "StHbtFemtoDstReader[ERROR]: cannot allocate memory for TSystem\n";
    return 0;
  }
  void *lDir = gSystem->OpenDirectory(aDir);
  if (!lDir) {
    std::cout << "StHbtFemtoDstReader[ERROR]: can't open directory " 
	      << aDir << std::endl;
    delete gSystem;
    return 0;
  }
  
  int mCount = 0;
  char *lFileName;

  while ((lFileName = (char *)gSystem->GetDirEntry(aDir))) {
    if (strcmp(lFileName, ".") == 0 || strcmp(lFileName, "..") == 0)
      continue;
    if (strstr(lFileName, aFilter)) {
      char *lFullName = gSystem->ConcatFileName(aDir, lFileName);
      std::cout << "StHbtFemtoDstReader[INFO]: Adding "
		<< lFullName << " to the chain" << std::endl;
      aChain->Add(lFullName);
      delete lFullName;
      mCount++;
      if ((aMaxFiles > 0) && (mCount > aMaxFiles)) break;
    }
  }
  
  std::cout << "StHbtFemtoDstReader[INFO]: Added " << mCount 
	    << " files to the chain" << std::endl;
  delete gSystem;
  return mCount;
}

//_________________
int StHbtFemtoDstReader::FillChain(TChain *aChain, const char *aFileName,
				   int aMaxFiles) {

  std::ifstream lInStream(aFileName);

  if (!lInStream.is_open()) {
    std::cout << "StHbtFemtoDstReader[ERROR]: can't open file list: "
	      << aFileName << std::endl;
    return 0;
  }
  std::cout << "StHbtFemtoDstReader[INFO]: using file list: "
	    << aFileName << std::endl;

  int mCount = 0;
  char buf[0xFF];
  for ( ; lInStream.good(); ) {
    lInStream.getline(buf, 0xFF);
    TString lFileName(buf);
    if (lFileName.Contains("root")) {
      aChain->Add(buf);
      mCount++;
      if ((aMaxFiles > 0) && (mCount > aMaxFiles)) break;
    }
  }

  std::cout << "StHbtFemtoDstReader[INFO]: Added " << mCount
	    << " files to the chain" << std::endl;
  return mCount;
}

//_________________
int StHbtFemtoDstReader::InitRead(string aDir, string aFileName,
				  string aFilter, int aMaxFiles) {

  mEventIndex = 0;
  //  mFemtoEvent = new StFemtoEvent();
  mTChain = new TChain("StFemtoDst", "StFemtoDst");
  
  std::cout << "StHbtFemtoDstReader[INFO]: Call InitRead" << std::endl;
  
  int lNumFiles = 0;
  if (!aFileName.empty()) {
    if (strstr(aFileName.c_str(), ".lis") ||
	strstr(aFileName.c_str(), ".list")) {
      lNumFiles = FillChain( mTChain, 
			     (aDir + aFileName).c_str(), 
			     aMaxFiles );
    }
    else {
      mTChain->Add((aDir + aFileName).c_str());
      lNumFiles++;
    }
  } else {
    lNumFiles = FillChain( mTChain, (char *)aDir.c_str(), // not cool
			   aFilter.c_str(), aMaxFiles );
  }
  mTChain->SetBranchAddress("StFemtoEvent", &mFemtoEvent);
  mNEvents = (unsigned int)mTChain->GetEntries();
  return lNumFiles;
}

//_________________
void StHbtFemtoDstReader::UninitRead() {
  //  if (mFemtoEvent) delete mFemtoEvent;
  if (mTChain) delete mTChain;
  mFemtoEvent = 0;
  mTChain   = 0;
}

//_________________
void StHbtFemtoDstReader::PrintTotalTracks() {

  std::cout << "StHbtFemtoDstReader[INFO]: mTotalTracks = "
	    << mTotalTracks << std::endl;
}

//_________________
int StHbtFemtoDstReader::GetNEvents() {

  if (mTChain) {
    return mNEvents;
  }
  else {
    return -1;
  }
}

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


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

  mFemtoEvent->Clear();
  int lBytes = mTChain->GetEntry(mEventIndex++);

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

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

  StHbtEvent *mHbtEvent = NULL;

  if(mFemtoEvent) {

    mHbtEvent = new StHbtEvent();

    //Check trigger Id
    if (!mTrigIdCol.empty()) {
      bool lTrigPass = false;
      unsigned int lTrigId = mFemtoEvent->GetTriggerId();
      if (((lTrigId & mTrig) == 0) || ((lTrigId & (!mTrig)) != 0))
        return mHbtEvent;
    }

    float mMagField = mFemtoEvent->GetMagneticField();
    StThreeVectorF v( mFemtoEvent->GetVertexPositionX(),
		      mFemtoEvent->GetVertexPositionY(),
		      mFemtoEvent->GetVertexPositionZ() );

    if(mDebug) {
      std::cout << Form("Vertex: x=%4.2f y=%4.2f z=%4.2f ",v.x(),v.y(),v.z()) << std::endl;
    }
    
    mHbtEvent->SetEventNumber( mFemtoEvent->GetEventId() );
    mHbtEvent->SetRunNumber( mFemtoEvent->GetRunId() );
    mHbtEvent->SetZdcAdcEast( mFemtoEvent->GetZDCe() );
    mHbtEvent->SetZdcAdcWest( mFemtoEvent->GetZDCw() );
    mHbtEvent->SetUncorrectedNumberOfPositivePrimaries( mFemtoEvent->GetRefMultPos() );
    mHbtEvent->SetUncorrectedNumberOfNegativePrimaries( mFemtoEvent->GetRefMultNeg() );
    mHbtEvent->SetUncorrectedNumberOfPrimaries( mFemtoEvent->GetRefMult() );
    mHbtEvent->SetPrimVertPos( v );
    mHbtEvent->SetMagneticField( mMagField );
    mHbtEvent->SetVpdVz( mFemtoEvent->GetVpdVz() );
    mHbtEvent->SetCent9( mFemtoEvent->GetCent9() );
    mHbtEvent->SetCent16( mFemtoEvent->GetCent16() );
    mHbtEvent->SetRefMultCorr( mFemtoEvent->GetRefMultCorr() );
    
    TClonesArray *mFemtoTrackArr = mFemtoEvent->GetFemtoTracks();
    
    if (mFemtoTrackArr) {
      int mNFemtoTracks = mFemtoTrackArr->GetEntries();
      mTotalTracks += (float)mNFemtoTracks;
      if (mDebug) {
	std::cout << "StHbtFemtoDstReader[DEBUG]: mNFemtoTracks = "
		  << mNFemtoTracks << std::endl;
      }
      
      for (int i = 0; i < mNFemtoTracks; i++) {

	StFemtoTrack *mFemtoTrack = (StFemtoTrack *)mFemtoTrackArr->At(i);
	StHbtTrack *mHbtTrack = new StHbtTrack;
	Short_t charge = mFemtoTrack->GetCharge();

	mHbtTrack->SetTrackId( mFemtoTrack->GetId() );
	//mHbtTrack->SetTrackType(mFemtoTrack->GetType());
	mHbtTrack->SetFlag( mFemtoTrack->GetFlag() );
	
	mHbtTrack->SetNHits( charge * mFemtoTrack->GetNHits() );
	//mHbtTrack->SetNHitsFit( mFemtoTrack->GetNHitsFit() );
	mHbtTrack->SetNHitsPossible( mFemtoTrack->GetNHitsPoss() );
	mHbtTrack->SetNHitsDedx( mFemtoTrack->GetNHitsDedx() );
	
	mHbtTrack->SetNSigmaElectron(mFemtoTrack->GetNSigmaElectron());
	mHbtTrack->SetNSigmaPion(mFemtoTrack->GetNSigmaPion());
	mHbtTrack->SetNSigmaKaon(mFemtoTrack->GetNSigmaKaon());
	mHbtTrack->SetNSigmaProton(mFemtoTrack->GetNSigmaProton());
	
	mHbtTrack->SetPidProbElectron(mFemtoTrack->GetPidProbElectron());
	mHbtTrack->SetPidProbPion(mFemtoTrack->GetPidProbPion());
	mHbtTrack->SetPidProbKaon(mFemtoTrack->GetPidProbKaon());
	mHbtTrack->SetPidProbProton(mFemtoTrack->GetPidProbProton());
	
	mHbtTrack->SetdEdx( mFemtoTrack->GetDedx() );
	mHbtTrack->SetChi2( mFemtoTrack->GetChi2() );
	
	mHbtTrack->SetDCAxyGlobal( mFemtoTrack->GetDCAxyGlobal() );
	mHbtTrack->SetDCAzGlobal( mFemtoTrack->GetDCAzGlobal() );


	//Primary track information
	StThreeVectorF mPrimMomentum( mFemtoTrack->GetPx(),
				      mFemtoTrack->GetPy(),
				      mFemtoTrack->GetPz() );
	StThreeVectorF mPrimOrigin( mFemtoEvent->GetVertexPositionX(),
				    mFemtoEvent->GetVertexPositionY(),
				    mFemtoEvent->GetVertexPositionZ() );
	StPhysicalHelixD mPrimHelix( mPrimMomentum, mPrimOrigin, 
				     mMagField * kilogauss, charge );

	//Global track information
	StThreeVectorF mGlobMomentum( mFemtoTrack->GetGlobalPx(),
				      mFemtoTrack->GetGlobalPy(),
				      mFemtoTrack->GetGlobalPz() );
	StThreeVectorF mGlobOrigin( (mFemtoTrack->GetDCAxGlobal() + 
				     mFemtoEvent->GetVertexPositionX()),
				    (mFemtoTrack->GetDCAyGlobal() + 
				     mFemtoEvent->GetVertexPositionY()),
				    (mFemtoTrack->GetDCAzGlobal() + 
				     mFemtoEvent->GetVertexPositionZ()) );
	StPhysicalHelixD mGlobHelix( mGlobMomentum, mGlobOrigin, 
				     mMagField * kilogauss, charge );

	//Set mP. Primary tracks have non-zero mP, globals have mP(0.,0.,0.)
	if( mFemtoTrack->GetType() == 0) { // Global tracks
	  mHbtTrack->SetP( StThreeVectorF(0., 0., 0.) );
	  mHbtTrack->SetHelix( mGlobHelix );
	  mHbtTrack->SetHelixGlobal( mGlobHelix );
	}
	else { //Primary tracks
	  mHbtTrack->SetP( mPrimMomentum );
	  mHbtTrack->SetHelix( mPrimHelix );
	  mHbtTrack->SetHelixGlobal( mGlobHelix );
	}
	//Set momentum of the global track
	mHbtTrack->SetPGlobal( mGlobMomentum );

	if(mDebug) {
	  std::cout << Form("Track parameters: px=%3.2f py=%3.2f pz=%3.2f q=%1d nhits=%2d",
			    mPrimMomentum.x(),mPrimMomentum.y(),mPrimMomentum.z(),
			    charge, mFemtoTrack->GetNHits() ) << std::endl;
	}
	
	mHbtTrack->SetTopologyMap( 0, mFemtoTrack->GetMap0() );
	mHbtTrack->SetTopologyMap( 1, mFemtoTrack->GetMap1() );
	
	mHbtTrack->SetTofBeta( mFemtoTrack->GetBeta() );
	mHbtTrack->SetHiddenInfo( 0 );
	
	mHbtEvent->TrackCollection()->push_back( mHbtTrack );
      } //for (int i = 0; i < mNFemtoTracks; i++)
    } //if (mFemtoTrackArr)
  } //if(mFemtoEvent)
  else {
    // FemtoDstMaker returned NULL, so we are not OK
    mReaderStatus = 1;
  }

  return mHbtEvent;
}

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

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

//_________________
StHbtString StHbtFemtoDstReader::Report() {

  StHbtString temp = "\nThis is the StHbtFemtoDstReader\n";
  return temp;
}


void StHbtFemtoDstReader::SetTrigger(unsigned int trig, char mode) {
  unsigned int trigColSize = mTrigIdCol.size();
  for (unsigned int iTrig = 0; iTrig < trigColSize; iTrig++) {
    if (mTrigIdCol[iTrig] == trig) {
      unsigned int mask = 1 << iTrig;
      if (!mode)
        mTrig = mTrig & (!mask);
      else
        mTrig = mTrig | mask;
      break;
    }
  }
  //
  // Output trigger mask
  //
  char bitMask[33] = {0};
  for (unsigned int iBit = 0; iBit < 32; iBit++) {
    bitMask[31 - iBit] = '0' + ((mTrig & (1 << iBit)) >> iBit);
  }

  std::cout << "StHbtFemtoDstReader[INFO]: Trigger mask = "
            << mTrig << " bit = "
            << bitMask << std::endl;
}

//_________________
void StHbtFemtoDstReader::AddTrigId(unsigned int aTrigId) {

  mTrigIdCol.push_back(aTrigId);
}

//_________________
void StHbtFemtoDstReader::DoDebug(bool aDebug) {
  mDebug = aDebug;
}
