/*
  Copyright (C) 2002-2022 CERN for the benefit of the ATLAS collaboration
*/

#ifndef XAOD_ANALYSIS
#include "TrkToolInterfaces/ITrackSelectorTool.h"
#include "TrkParametersIdentificationHelpers/TrackParametersIdHelper.h"

#include "xAODTau/TauJet.h"
#include "xAODTau/TauTrackContainer.h"

#include "TauTrackFinder.h"
#include "tauRecTools/TrackSort.h"


TauTrackFinder::TauTrackFinder(const std::string& name) :
  TauRecToolBase(name) {
  m_EMSamplings = {CaloSampling::EME1, CaloSampling::EMB1};
  m_HadSamplings = {CaloSampling::TileBar1, CaloSampling::HEC1, CaloSampling::TileExt1};
}

TauTrackFinder::~TauTrackFinder() {
}

// * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * *
StatusCode TauTrackFinder::initialize() {

  // retrieve tools
  ATH_CHECK( m_trackSelectorTool_tau.retrieve() );
  ATH_CHECK( m_trackToVertexTool.retrieve() );
  ATH_CHECK( m_caloExtensionTool.retrieve() );
  ATH_CHECK( m_trackToVertexIPEstimator.retrieve() );

  // initialize ReadHandleKey
  ATH_CHECK( m_trackPartInputContainer.initialize() );
  // use CaloExtensionTool when key is empty 
  ATH_CHECK( m_ParticleCacheKey.initialize(SG::AllowEmpty) );
  // allow empty for LRT
  ATH_CHECK( m_largeD0TracksInputContainer.initialize(SG::AllowEmpty) );

  if(m_useGhostTracks) {
    if(inTrigger()) {
      ATH_MSG_ERROR ("Ghost matching is not a valid tau-track association scheme for trigger, use cone association. Aborting.");
      return StatusCode::FAILURE;
    }
    ATH_MSG_INFO ("Using ghost matching for tau-track association" << m_ghostTrackDR );
    // allow empty for trigger 
    ATH_CHECK( m_jetContainer.initialize(SG::AllowEmpty) );
  }

  ATH_CHECK( m_beamSpotKey.initialize(inTrigger()) );

  return StatusCode::SUCCESS;
}

// * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * *
StatusCode TauTrackFinder::executeTrackFinder(xAOD::TauJet& pTau, xAOD::TauTrackContainer& tauTrackCon) const {
  
  std::vector<const xAOD::TrackParticle*> tauTracks;
  std::vector<const xAOD::TrackParticle*> wideTracks;
  std::vector<const xAOD::TrackParticle*> otherTracks;

  //Retrieve standard track container
  const xAOD::TrackParticleContainer* trackParticleCont = nullptr; 
  
  SG::ReadHandle<xAOD::TrackParticleContainer> trackPartInHandle( m_trackPartInputContainer );
  if (!trackPartInHandle.isValid()) {
    ATH_MSG_ERROR ("Could not retrieve HiveDataObj with key " << trackPartInHandle.key());
    return StatusCode::FAILURE;
  }
  trackParticleCont = trackPartInHandle.cptr();

  //Retrieve LRT container
  const xAOD::TrackParticleContainer* largeD0TracksParticleCont = nullptr; 
  
  if (! m_largeD0TracksInputContainer.empty()) { 
    SG::ReadHandle<xAOD::TrackParticleContainer> trackPartInHandle( m_largeD0TracksInputContainer );
    if (!trackPartInHandle.isValid()) {
      ATH_MSG_VERBOSE ("Could not retrieve HiveDataObj with key " << trackPartInHandle.key());
      ATH_MSG_VERBOSE ("LRT container " << trackPartInHandle.key()<<" is not being used for tau tracks");
    }
    else { 
      largeD0TracksParticleCont = trackPartInHandle.cptr();
    }  
  }

  // retrieve the seed jet container when using ghost-matching
  const xAOD::JetContainer* jetContainer = nullptr; 
  if (! m_jetContainer.empty()) {
    SG::ReadHandle<xAOD::JetContainer> jetContHandle( m_jetContainer );
    if (!jetContHandle.isValid()) {
      ATH_MSG_ERROR ("Could not retrieve HiveDataObj with key " << jetContHandle.key());
      return StatusCode::FAILURE;
    }
    jetContainer = jetContHandle.cptr();
  }

  // get the primary vertex
  const xAOD::Vertex* pVertex = pTau.vertex();

  // retrieve tracks wrt a vertex                                                                                                                              
  // as a vertex is used: tau origin / PV / beamspot / 0,0,0 (in this order, depending on availability)                                                        

  getTauTracksFromPV(pTau, *trackParticleCont, pVertex, m_useGhostTracks, jetContainer, tauTracks, wideTracks, otherTracks);

  bool foundLRTCont = bool (largeD0TracksParticleCont != nullptr);
  // additional LRT with vertex association added to tracks
  if (foundLRTCont){
    // for now, use cone association for LRTs, not ghost association
    getTauTracksFromPV(pTau, *largeD0TracksParticleCont, pVertex, false, nullptr, tauTracks, wideTracks, otherTracks);
  }

  // remove core and wide tracks outside a maximal delta z0 wrt lead core track                                                                                
  if (m_applyZ0cut) {
    this->removeOffsideTracksWrtLeadTrk(tauTracks, wideTracks, otherTracks, pVertex, m_z0maxDelta);
  }

  if(m_removeDuplicateCoreTracks){
    bool alreadyUsed = false;
    for (std::vector<const xAOD::TrackParticle*>::iterator track_it = tauTracks.begin(); track_it != tauTracks.end() ;)
      {
	alreadyUsed = false;
	//loop over all up-to-now core tracks	
	for( const xAOD::TauTrack* tau_trk : tauTrackCon ) {
	  //originally it was coreTrack&passTrkSelector
	  if(! tau_trk->flagWithMask( (1<<xAOD::TauJetParameters::TauTrackFlag::coreTrack) | (1<<xAOD::TauJetParameters::TauTrackFlag::passTrkSelector))) continue; 
	  if( (*track_it) == tau_trk->track()) alreadyUsed = true;
	}
	//if this track has already been used by another tau, don't associate it to this new one                                                               
	if(alreadyUsed) ATH_MSG_INFO( "Found Already Used track new, now removing: " << *track_it );
	if (alreadyUsed) track_it = tauTracks.erase(track_it);
	else ++track_it;
      }
  }

  // associate track to tau candidate and calculate charge                                                                                                    
  float charge = 0.;  
  for (unsigned int i = 0; i < tauTracks.size(); ++i) {
    const xAOD::TrackParticle* trackParticle = tauTracks.at(i);

    ATH_MSG_VERBOSE(name() << " adding core track nr: " << i
		    << " eta " << trackParticle->eta()
		    << " phi " << trackParticle->phi());

    charge += trackParticle->charge();

    xAOD::TauTrack* track = new xAOD::TauTrack();
    tauTrackCon.push_back(track);

    ElementLink<xAOD::TrackParticleContainer> linkToTrackParticle;
    if (foundLRTCont && isLargeD0Track(trackParticle)){//Check LRT track and link to container
      linkToTrackParticle.toContainedElement(*largeD0TracksParticleCont, trackParticle);
      track->setFlag(xAOD::TauJetParameters::TauTrackFlag::LargeRadiusTrack, true);
    }
    else {
      linkToTrackParticle.toContainedElement(*trackParticleCont, trackParticle);
      track->setFlag(xAOD::TauJetParameters::TauTrackFlag::LargeRadiusTrack, false);
    }
    track->addTrackLink(linkToTrackParticle);

    track->setP4(trackParticle->pt(), trackParticle->eta(), trackParticle->phi(), trackParticle->m());
    track->setFlag(xAOD::TauJetParameters::TauTrackFlag::coreTrack, true);
    track->setFlag(xAOD::TauJetParameters::TauTrackFlag::passTrkSelector, true);
    // in case TrackClassifier is not run, still get sensible results                                                                                        
    track->setFlag(xAOD::TauJetParameters::TauTrackFlag::classifiedCharged, true);
    track->setFlag(xAOD::TauJetParameters::TauTrackFlag::unclassified, true);
    
    ElementLink<xAOD::TauTrackContainer> linkToTauTrack;
    linkToTauTrack.toContainedElement(tauTrackCon, track);
    pTau.addTauTrackLink(linkToTauTrack);

    ATH_MSG_VERBOSE(name() << " added core track nr: " << i
		    << " eta " << pTau.track(i)->eta()
		    << " phi " << pTau.track(i)->phi());
  }
  // set the charge, which is defined by the core tau tracks only                                                                                              
  pTau.setCharge(charge);

  for (unsigned int i = 0; i < wideTracks.size(); ++i) {
    const xAOD::TrackParticle* trackParticle = wideTracks.at(i);

    ATH_MSG_VERBOSE(name() << " adding wide track nr: " << i
		    << " eta " << trackParticle->eta()
		    << " phi " << trackParticle->phi());

    xAOD::TauTrack* track = new xAOD::TauTrack();
    tauTrackCon.push_back(track);

    ElementLink<xAOD::TrackParticleContainer> linkToTrackParticle;
    if (foundLRTCont && isLargeD0Track(trackParticle)){//Check LRT track and link to container
      linkToTrackParticle.toContainedElement(*largeD0TracksParticleCont, trackParticle);
      track->setFlag(xAOD::TauJetParameters::TauTrackFlag::LargeRadiusTrack, true);
    }
    else {
      linkToTrackParticle.toContainedElement(*trackParticleCont, trackParticle);
      track->setFlag(xAOD::TauJetParameters::TauTrackFlag::LargeRadiusTrack, false);
    }
    track->addTrackLink(linkToTrackParticle);

    track->setP4(trackParticle->pt(), trackParticle->eta(), trackParticle->phi(), trackParticle->m());
    track->setFlag(xAOD::TauJetParameters::TauTrackFlag::wideTrack, true);
    track->setFlag(xAOD::TauJetParameters::TauTrackFlag::passTrkSelector, true);
    // in case TrackClassifier is not run, still get sensible results                                                                                        
    track->setFlag(xAOD::TauJetParameters::TauTrackFlag::classifiedIsolation, true); // for sake of trigger, reset in TauTrackClassifier
    track->setFlag(xAOD::TauJetParameters::TauTrackFlag::modifiedIsolationTrack, true); // for sake of trigger, reset in TauTrackClassifier
    track->setFlag(xAOD::TauJetParameters::TauTrackFlag::unclassified, true);

    ElementLink<xAOD::TauTrackContainer> linkToTauTrack;
    linkToTauTrack.toContainedElement(tauTrackCon, track);
    pTau.addTauTrackLink(linkToTauTrack);
  }

  //These are set again in TauTrackClassifier                                                                                                                  
  pTau.setDetail(xAOD::TauJetParameters::nChargedTracks, (int) pTau.nTracks());
  pTau.setDetail(xAOD::TauJetParameters::nIsolatedTracks, (int) pTau.nTracks(xAOD::TauJetParameters::classifiedIsolation));

  for (unsigned int i = 0; i < otherTracks.size(); ++i) {
    const xAOD::TrackParticle* trackParticle = otherTracks.at(i);

    ATH_MSG_VERBOSE(name()     << " adding other track nr: " << i
		    << " eta " << trackParticle->eta()
		    << " phi " << trackParticle->phi());

    xAOD::TauTrack* track = new xAOD::TauTrack();
    tauTrackCon.push_back(track);

    ElementLink<xAOD::TrackParticleContainer> linkToTrackParticle;
    if (foundLRTCont && isLargeD0Track(trackParticle)){//Check LRT track and link to container
      linkToTrackParticle.toContainedElement(*largeD0TracksParticleCont, trackParticle);
      track->setFlag(xAOD::TauJetParameters::TauTrackFlag::LargeRadiusTrack, true); 
    }
    else {
      linkToTrackParticle.toContainedElement(*trackParticleCont, trackParticle);
      track->setFlag(xAOD::TauJetParameters::TauTrackFlag::LargeRadiusTrack, false);
    }
    track->addTrackLink(linkToTrackParticle);

    track->setP4(trackParticle->pt(), trackParticle->eta(), trackParticle->phi(), trackParticle->m());
    float dR = track->p4().DeltaR(pTau.p4());
    if(dR<=0.2) track->setFlag(xAOD::TauJetParameters::TauTrackFlag::coreTrack, true);
    else track->setFlag(xAOD::TauJetParameters::TauTrackFlag::wideTrack, true);
    track->setFlag(xAOD::TauJetParameters::TauTrackFlag::unclassified, true);

    ElementLink<xAOD::TauTrackContainer> linkToTauTrack;
    linkToTauTrack.toContainedElement(tauTrackCon, track);
    pTau.addTauTrackLink(linkToTauTrack);
  }

  pTau.setDetail(xAOD::TauJetParameters::nLargeRadiusTracks, (int) pTau.nTracks(xAOD::TauJetParameters::LargeRadiusTrack));
  // keep track of total number of associated tracks, in case of tau track thinning
  pTau.setDetail(xAOD::TauJetParameters::nAllTracks, (int) pTau.nAllTracks());

  ATH_MSG_DEBUG("numTrack: " << "/" << pTau.nTracks());
  ATH_MSG_DEBUG("charge: " << "/" << pTau.charge());

  // impact parameter variables w.r.t. tau vertex 
  const xAOD::Vertex* vxcand = nullptr;

  xAOD::Vertex vxbkp;
  vxbkp.makePrivateStore();

  // FIXME: don't we want to use the beamspot in the offline reconstruction too, when there is no reconstructed primary vertex?
  if (pTau.vertex()!=nullptr && pTau.vertex()->vertexType() != xAOD::VxType::NoVtx) {
    vxcand = pTau.vertex();
  }
  else if (inTrigger()) { // online: use vertex with x-y coordinates from the beamspot and the z from the leading track
    vxbkp.setX(0); vxbkp.setY(0); vxbkp.setZ(0);

    SG::ReadCondHandle<InDet::BeamSpotData> beamSpotHandle { m_beamSpotKey };
    if(beamSpotHandle.isValid()) {
      vxbkp.setPosition(beamSpotHandle->beamPos());
      const auto& cov = beamSpotHandle->beamVtx().covariancePosition();
      vxbkp.setCovariancePosition(cov);

      if(!tauTracks.empty()) {
         vxbkp.setZ(tauTracks.at(0)->z0());
      }
    }
    else {
      ATH_MSG_DEBUG("No Beamspot object in tau candidate");
    }
    vxcand = & vxbkp;
  }


  // this could be replaced with TauTrack::setDetail
  static const SG::AuxElement::Accessor<float> dec_d0TJVA("d0TJVA");
  static const SG::AuxElement::Accessor<float> dec_z0sinthetaTJVA("z0sinthetaTJVA");
  static const SG::AuxElement::Accessor<float> dec_d0SigTJVA("d0SigTJVA");
  static const SG::AuxElement::Accessor<float> dec_z0sinthetaSigTJVA("z0sinthetaSigTJVA");

  for(auto *track : pTau.allTracks()) {      
    dec_d0TJVA(*track) = track->track()->d0();
    dec_z0sinthetaTJVA(*track) = track->z0sinThetaTJVA(pTau);
    dec_d0SigTJVA(*track) = -999.;
    dec_z0sinthetaSigTJVA(*track) = -999.;

    // in the trigger, z0sintheta and corresponding significance are meaningless if we use the beamspot
    if(vxcand) {
      std::unique_ptr<const Trk::ImpactParametersAndSigma> myIPandSigma 
	= std::unique_ptr<const Trk::ImpactParametersAndSigma>(m_trackToVertexIPEstimator->estimate(track->track(), vxcand));
      
      if(myIPandSigma) {
	dec_d0TJVA(*track) = myIPandSigma->IPd0;
	dec_z0sinthetaTJVA(*track) = myIPandSigma->IPz0SinTheta;
	dec_d0SigTJVA(*track) = (myIPandSigma->sigmad0 != 0.) ? (float)( myIPandSigma->IPd0 / myIPandSigma->sigmad0 ) : -999.;
	dec_z0sinthetaSigTJVA(*track) = (myIPandSigma->sigmaz0SinTheta != 0.) ? (float)( myIPandSigma->IPz0SinTheta / myIPandSigma->sigmaz0SinTheta ) : -999.;
      }
    }
  }

  // extrapolate core tracks to calorimeter surface
  // store information only in ExtraDetailsContainer
  if(!m_bypassExtrapolator)
    {
      StatusCode sc = extrapolateToCaloSurface(pTau);
      if (sc.isFailure() && !sc.isRecoverable()) {
	ATH_MSG_ERROR("couldn't extrapolate tracks to calo surface");
	return StatusCode::FAILURE;
      }
    }
  
  return StatusCode::SUCCESS;
}


// * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * *
TauTrackFinder::TauTrackType TauTrackFinder::tauTrackType( const xAOD::TauJet& pTau,
							   const xAOD::TrackParticle& trackParticle,
							   const xAOD::Vertex* primaryVertex) const
{
  double dR = pTau.p4().DeltaR(trackParticle.p4());

  if (dR > m_maxJetDr_wide) return NotTauTrack;

  bool goodTrack = true;
  if(!m_bypassSelector)
    goodTrack = m_trackSelectorTool_tau->decision(trackParticle, primaryVertex);
    
  if (goodTrack) {
    if (dR > m_maxJetDr_tau)
      return TauTrackWide;
    else
      return TauTrackCore;
  } else
    return TauTrackOther;
}

// * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * *
void TauTrackFinder::getTauTracksFromPV( const xAOD::TauJet& pTau,
					 const xAOD::TrackParticleContainer& trackParticleCont,
					 const xAOD::Vertex* primaryVertex,
					 const bool& useGhostTracks,
					 const xAOD::JetContainer* jetContainer,
					 std::vector<const xAOD::TrackParticle*> &tauTracks,
					 std::vector<const xAOD::TrackParticle*> &wideTracks,
					 std::vector<const xAOD::TrackParticle*> &otherTracks) const
{
  std::vector<const xAOD::TrackParticle*> ghostTracks;
  if(useGhostTracks) {
    const xAOD::Jet* seedJet = pTau.jet();
    if(seedJet) {
      if(!seedJet->getAssociatedObjects("GhostTrack", ghostTracks)) {
	ATH_MSG_WARNING("Could not retrieve GhostTrack from seed jet.");
      }
    }
  }
  
  for (const xAOD::TrackParticle *trackParticle : trackParticleCont) {
    TauTrackType type = tauTrackType(pTau, *trackParticle, primaryVertex);
    if(type == NotTauTrack) continue;

    if(useGhostTracks) {
      // require that tracks are ghost-matched with the seed jet at large dR(tau,track), to avoid using tracks from another tau
      double dR = pTau.p4().DeltaR(trackParticle->p4());
      if (dR > m_ghostTrackDR) {
	if (std::find(ghostTracks.begin(), ghostTracks.end(), trackParticle) == ghostTracks.end()) {
	  // check whether the jet closest to the track is the current seed jet
	  // if so, recover the track even if not ghost-matched, to improve tau-track association efficiency at low pt (esp. for 3p)
	  double dRmin = 999.;
	  bool isSeedClosest = false;
	  for (const xAOD::Jet* jet : *jetContainer) {
	    xAOD::JetFourMom_t jetP4 = jet->jetP4(xAOD::JetScale::JetConstitScaleMomentum);
	    TLorentzVector jetLV;
	    jetLV.SetPtEtaPhiM(jetP4.Pt(), jetP4.Eta(), jetP4.Phi(), jetP4.M());
	    double dRjet = trackParticle->p4().DeltaR(jetLV);
	    if(dRjet < dRmin) {
	      dRmin = dRjet;
	      isSeedClosest = (jet == pTau.jet());
	    }
	  }
	  if(!isSeedClosest) continue;
	}
      }
    }
    
    if (type == TauTrackCore)
      tauTracks.push_back(trackParticle);
    else if (type == TauTrackWide)
      wideTracks.push_back(trackParticle);
    else if (type == TauTrackOther)
      otherTracks.push_back(trackParticle);
  }
  std::sort(tauTracks.begin(), tauTracks.end(), TrackSort());
  std::sort(wideTracks.begin(), wideTracks.end(), TrackSort());
  std::sort(otherTracks.begin(), otherTracks.end(), TrackSort());
}


// * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * *
StatusCode TauTrackFinder::extrapolateToCaloSurface(xAOD::TauJet& pTau) const {

  Trk::TrackParametersIdHelper parsIdHelper;

  int trackIndex = -1;
  const Trk::CaloExtension * caloExtension = nullptr;
  std::unique_ptr<Trk::CaloExtension> uniqueExtension;
  for( xAOD::TauTrack* tauTrack : pTau.allTracks() ) {
    const xAOD::TrackParticle *orgTrack = tauTrack->track();
    if( !orgTrack ) continue;
    trackIndex = orgTrack->index();   

    // set default values
    float etaEM = -10.0;
    float phiEM = -10.0;
    float etaHad = -10.0;
    float phiHad = -10.0;

    // get the extrapolation into the calo
    ATH_MSG_DEBUG( "Try extrapolation of track with pt = " << orgTrack->pt() 
		   << ", eta " << orgTrack->eta() 
		   << ", phi" << orgTrack->phi() );

    if(!m_ParticleCacheKey.key().empty()){
      /*get the CaloExtension object*/
      ATH_MSG_VERBOSE("Using the CaloExtensionBuilder Cache");
      SG::ReadHandle<CaloExtensionCollection>  particleCache {m_ParticleCacheKey};
      caloExtension = (*particleCache)[trackIndex];
      ATH_MSG_VERBOSE("Getting element " << trackIndex << " from the particleCache");
      if( not caloExtension ){
        ATH_MSG_VERBOSE("Cache does not contain a calo extension -> "
                        "Calculating with the a CaloExtensionTool");
        uniqueExtension = m_caloExtensionTool->caloExtension(
          Gaudi::Hive::currentContext(), *orgTrack);
        caloExtension = uniqueExtension.get();
      }
    }
    else {
      /* If CaloExtensionBuilder is unavailable, use the calo extension tool */
      ATH_MSG_VERBOSE("Using the CaloExtensionTool");
      uniqueExtension = m_caloExtensionTool->caloExtension(
        Gaudi::Hive::currentContext(), *orgTrack);
      caloExtension = uniqueExtension.get();
    }

    if (!caloExtension)
      { 
	ATH_MSG_DEBUG("Track extrapolation failed");
      }
    else {
      const std::vector<Trk::CurvilinearParameters>& clParametersVector = caloExtension->caloLayerIntersections();
      if (clParametersVector.empty()) {
	ATH_MSG_DEBUG("Track extrapolation failed");
      }

      ATH_MSG_DEBUG("Scanning samplings");
      bool validECal = false;
      bool validHCal = false;
      for( const Trk::CurvilinearParameters& cur : clParametersVector ){
	ATH_MSG_DEBUG("Sampling " << parsIdHelper.caloSample(cur.cIdentifier()) );
                
	// only use entry layer
	if( not parsIdHelper.isEntryToVolume(cur.cIdentifier()) ) continue;

	CaloSampling::CaloSample sample = parsIdHelper.caloSample(cur.cIdentifier());

	// ECal
	if( not validECal and m_EMSamplings.count(sample))
	  {
	    validECal = true;
	    etaEM = cur.position().eta();
	    phiEM = cur.position().phi();
	    ATH_MSG_DEBUG("Extrapolated to ECal layer " << sample);
	  }

	// HCal
	if( not validHCal and m_HadSamplings.count(sample))
	  {
	    validHCal = true;
	    etaHad = cur.position().eta();
	    phiHad = cur.position().phi();
	    ATH_MSG_DEBUG("Extrapolated to HCal layer " << sample);
	  }
	if( validECal and validHCal ) break;
      }
      // EM failure warn if within acceptance 
      if( not validECal and std::abs(orgTrack->pt()) < 2.48 ){
	ATH_MSG_DEBUG("Failed extrapolation to ECal");
      }
      // Had failure warn if enough pt to reach HCal
      if( not validHCal and orgTrack->pt() > 2000. ){
	ATH_MSG_DEBUG("Failed extrapolation to HCal");
      }

      ATH_MSG_DEBUG( "Extrapolated track with eta=" << orgTrack->eta()
		     << " phi="<<orgTrack->phi()
		     << " to ECal eta=" << etaEM 
		     << " phi="<< phiEM
		     << " HCal eta=" << etaHad 
		     << " phi="<< phiHad
		     );
    }
    tauTrack->setDetail(xAOD::TauJetParameters::CaloSamplingEtaEM, etaEM);
    tauTrack->setDetail(xAOD::TauJetParameters::CaloSamplingPhiEM, phiEM);
    tauTrack->setDetail(xAOD::TauJetParameters::CaloSamplingEtaHad, etaHad);
    tauTrack->setDetail(xAOD::TauJetParameters::CaloSamplingPhiHad, phiHad);
  }

  return StatusCode::SUCCESS;

}

// * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * *
void TauTrackFinder::removeOffsideTracksWrtLeadTrk(std::vector<const xAOD::TrackParticle*> &tauTracks,
						   std::vector<const xAOD::TrackParticle*> &wideTracks,
						   std::vector<const xAOD::TrackParticle*> &otherTracks,
						   const xAOD::Vertex* tauOrigin,
						   double maxDeltaZ0) const
{
  float MAX=1e5;

  // need at least one core track to have a leading trk to compare with
  if (tauTracks.empty()) return;

  // get lead trk parameters
  const xAOD::TrackParticle *leadTrack = tauTracks.at(0);
  float z0_leadTrk = getZ0(leadTrack, tauOrigin);

  if (z0_leadTrk > MAX-1) return; // bad lead trk -> do nothing

  ATH_MSG_VERBOSE("before z0 cut: #coreTracks=" << tauTracks.size() << ", #wideTracks=" << wideTracks.size() << ", #otherTracks=" << otherTracks.size());

  std::vector<const xAOD::TrackParticle*>::iterator itr;
    
  // skip leading track, because it is the reference
  itr = tauTracks.begin()+1;
  while (itr!=tauTracks.end()) {
    float z0 = getZ0(*itr, tauOrigin);
    float deltaZ0=z0 - z0_leadTrk;
    ATH_MSG_VERBOSE("core Trks: deltaZ0= " << deltaZ0);

    if ( std::abs(deltaZ0) < maxDeltaZ0 ) {++itr;}
    else {
      if (m_storeInOtherTrks) otherTracks.push_back(*itr);
      itr = tauTracks.erase(itr); //remove from core track collection
    }
  }

  // check wide tracks
  itr = wideTracks.begin();
  while (itr!=wideTracks.end()) {
    float z0 = getZ0(*itr, tauOrigin);
    float deltaZ0=z0 - z0_leadTrk;
    ATH_MSG_VERBOSE("wide Trks: deltaZ0= " << deltaZ0);

    if ( std::abs(deltaZ0) < maxDeltaZ0 ) { ++itr; }
    else {
      if (m_storeInOtherTrks) otherTracks.push_back(*itr);
      itr = wideTracks.erase(itr); //remove from wide track collection
    }
  }

  ATH_MSG_VERBOSE("after z0 cut: #coreTracks=" << tauTracks.size() << ", #wideTracks=" << wideTracks.size() << ", #otherTracks=" << otherTracks.size());

  // sort again
  std::sort(tauTracks.begin(), tauTracks.end(), TrackSort());
  std::sort(wideTracks.begin(), wideTracks.end(), TrackSort());
  std::sort(otherTracks.begin(), otherTracks.end(), TrackSort());
}

// * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * *
float TauTrackFinder::getZ0(const xAOD::TrackParticle* track, const xAOD::Vertex* vertex) const
{
  float MAX=1e5;

  if (!track) return MAX;

  const Trk::Perigee* perigee = nullptr;
  if (vertex) perigee = m_trackToVertexTool->perigeeAtVertex(*track, vertex->position());
  else        perigee = m_trackToVertexTool->perigeeAtVertex(*track); //will use beamspot or 0,0,0 instead

  if (!perigee) {
    ATH_MSG_WARNING("Bad track; can't find perigee at vertex.");
    return MAX;
  }

  float z0 = perigee->parameters()[Trk::z0];

  delete perigee; //cleanup necessary to prevent mem leak

  return z0;
}

bool TauTrackFinder::isLargeD0Track(const xAOD::TrackParticle* track) const
{
  const std::bitset<xAOD::NumberOfTrackRecoInfo> patternReco = track->patternRecoInfo();
  if (patternReco.test(xAOD::TrackPatternRecoInfo::SiSpacePointsSeedMaker_LargeD0)) {
    ATH_MSG_DEBUG("LargeD0Track found");
    return true;
  }

  return false;
}
#endif