////////////////////////////////////////////////////////////////////
//   File:  TauVariables.cc
//   Authors: Matt Reece
//   Purpose: High level analysis tau variables defined as functions
//   Created: 04/01/2002
//
//   Revision:  0.1
//
////////////////////////////////////////////////////////////////////

#include "HighLevelObjects/TauVariables.hh"
#include "HighLevelObjects/JetVariables.hh"
#include "HighLevelObjects/TrackingVariables.hh"
#include "TauObjects/CdfTau.hh"
#include <math.h>
#include <algorithm>

namespace TauVariables {
  int tauNumTowers(const CdfTau& tau) {
    return tau.nTowers();
  }
  int tauNumClumpTowers(const CdfTau& tau) {
    return tau.nClumpTowers();
  }

  int tauNumChargedTracks(const CdfTau& tau) {
    return tau.chargedTracksNumber();
  }

  int tauNumChargedTracks10deg(const CdfTau& tau) {
    return tau.chargedTracksNumberSeed10Degree();
  }

  int tauNumChargedTracks30deg(const CdfTau& tau) {
    return tau.chargedTracksNumberSeed30Degree();
  }

  int tauNumChargedTracks10to30deg(const CdfTau& tau) {
    return tau.chargedTracksNumberSeed10To30Degree();
  }

  int tauChargeIn10deg(const CdfTau& tau) {
    return tau.tracksChargeInSeed10Degree();
  }

  int tauNumWrongVtxTrk10deg(const CdfTau& tau) {
    return tau.wrongVertexTracksNumberSeed10Degree();
  }

  double tauEmFraction(const CdfTau& tau) {
    return tau.emFraction();
  }

  double tauClusterDetectorEta(const CdfTau& tau) {
    return tau.clusterDetectorEta();
  }

  double tauSeedTrkDetectorEta(const CdfTau& tau) {
    return tau.seedTrackDetectorEta();
  }

  double tauEtaEta(const CdfTau& tau) {
    return tau.clusterEtaEtaMoment();
  }

  double tauPhiPhi(const CdfTau& tau) {
    return tau.clusterPhiPhiMoment();
  }

  double tauDelr(const CdfTau& tau) { 
    double ee = tau.clusterEtaEtaMoment(), pp = tau.clusterPhiPhiMoment();
    return sqrt(ee*ee+pp*pp);
  }

  double tauEmEnergy(const CdfTau& tau) {
    return tau.clusterEmE();
  }

  double tauHadEnergy(const CdfTau& tau) {
    return tau.clusterHadronE();
  }

  double tauEmEt(const CdfTau& tau) {
    return tau.clusterEmEt();
  }

  double tauHadEt(const CdfTau& tau) {
    return tau.clusterHadronEt();
  }

  double tauEt(const CdfTau& tau) {
    return tau.clusterEmEt()+tau.clusterHadronEt();
  }

  double tauCaloEtIso(const CdfTau& tau) {
    return tau.caloIsolationEt();  
  }

  double tauCaloHadEtIso(const CdfTau& tau) {
    return tau.caloIsolationHadronEt();
  }

  double tauCaloEmEtIso(const CdfTau& tau) {
    return tau.caloIsolationEmEt();
  }

  double tauReferenceVtxZ(const CdfTau& tau) {
    return tau.referenceVertexZ();
  }

  double tauClusterEta(const CdfTau& tau) {
    double theta = tauClusterMomentum(tau).theta();
    return -log(tan(theta/2.0));
  }

  double tauClusterPhi(const CdfTau& tau) {
    return tauClusterMomentum(tau).phi();
  }

  HepLorentzVector tauClusterMomentum(const CdfTau& tau) {
    return HepLorentzVector(tau.clusterPx(), tau.clusterPy(),
                            tau.clusterPz(), tau.clusterE());
  }

  HepLorentzVector tauTrackMomentum(const CdfTau& tau) {
    // Here's one way of trying it... alternatively we could
    // get the actual track collection and use that.
    double pt = tauTracksPt(tau);
    double eta = tauTracksEta(tau);
    double phi = tauTracksPhi(tau);
    double mass = tauTracksMass(tau);
    double px = pt*cos(phi);
    double py = pt*sin(phi);
    double pz = pt * sinh(eta); // cot(theta) = sinh(eta)
    double E = sqrt(px*px+py*py+pz*pz+mass*mass);
    return HepLorentzVector(px, py, pz, E);
  }

  HepLorentzVector tauVisMomentum(const CdfTau& tau) {
    HepLorentzVector v;
    tau.pi0sFourMomentum(&v); // Fill with sum of pi0 four-momentum
    v += tauTrackMomentum(tau);

    return v;
  }

  double tauTracksPt(const CdfTau& tau) {
    return tau.tracksPt();
  }

  double tauTracksScalarPtSum(const CdfTau& tau) {
    return tau.tracksScalarPtSum();
  }

  double tauTracksMass(const CdfTau& tau) {
    return tau.tracksMass();
  }

  double tauTracksEta(const CdfTau& tau) {
    return tau.tracksEta();
  }

  double tauTracksPhi(const CdfTau& tau) {
    return tau.tracksPhi();
  }

  double tauTrackIso(const CdfTau& tau) {
    return tau.trackIsolation();
  }

  double tauEp(const CdfTau& tau) {
    double Et = tauEt(tau);
    double pt = tauTracksPt(tau);
    if(pt != 0) {
      return Et / pt;
    }
    return -9999.0;
  }

  double tauEpvis(const CdfTau& tau) {
    double Et = tauEt(tau);
    double pt = tauPi0AndTrackPt(tau);
    if(pt != 0) {
      return Et / pt;
    }
    return -9999.0;
  }

  double tauSeedTrackVtxZ(const CdfTau& tau) {
    return tau.seedTrackVertexZ();
  }

  double tauSeedTrackPt(const CdfTau& tau) {
    return tau.seedTrackPt();
  }

  double tauSeedTrackDetectorEta(const CdfTau& tau) {
    // The function in CdfTau seems to always return 0,
    // so let's mimic the FlatNtuple extreta code....
    const CdfTrack * seed = tau.seedTrack();
    if(fabs(TrackingVariables::trkZ0(seed)) > TLZPESL0) return -10;
    const HepPoint3D& hpt = TrackingVariables::trkExtPos(TLYCES, seed);
    if(hpt.z() == -999999.) return -10;
    if(fabs(hpt.z()) < TLZPESL0) return hpt.eta();
    // else we got to the PES radius
    HepPlane3D plug(0,0,1,(TrackingVariables::trkEta0(seed) > 0 ?
                           -(TLZPESL0+TLZPESL1)/2 :
                           (TLZPESL0+TLZPESL1)/2));
    const HepPoint3D& hpt2 = 
      seed->getTrajectory().newIntersectionWith(plug)->position();
    if(hpt2.perp() > 0) return hpt2.eta();
    return -10;
  }

  double tauSeedTrackPhi(const CdfTau& tau) {
    Hep3Vector seed = tau.seedTrack()->momentum();
    return seed.phi();
  }

  double tauAngleSeedToCluster(const CdfTau& tau) {
    return tau.angleSeedTrackToCluster();
  }

  double tauPhiSeedToCluster(const CdfTau& tau) {
    return tau.phiAngleSeedTrackToCluster();
  }

  double tauEtaSeedToCluster(const CdfTau& tau) {
    return tau.etaAngleSeedTrackToCluster();
  }

  int tauNumPi0s(const CdfTau& tau) {
    return tau.pi0sNumber();
  }

  double tauPi0AndTrackMass(const CdfTau& tau) {
    return tau.pi0sAndTracksMass();
  }

  double tauPi0AndTrackPt(const CdfTau& tau) {
    return tau.pi0sAndTracksPt();
  }

  int tauNumMuonStubs(const CdfTau& tau) {
    return tau.muonStubsNumber();
  }

  int tauNumMuonHits(const CdfTau& tau) {
    return tau.muonHitsNumber();
  }

  double pi0Eta(const Pi0Candidate& pi0) {
    return pi0.eta();
  }

  double pi0Phi(const Pi0Candidate& pi0) {
    return pi0.phi();
  }

  double pi0Energy(const Pi0Candidate& pi0) {
    return pi0.energy();
  }

  double pi0DetectorZ(const Pi0Candidate& pi0) {
    return pi0.detectorZ();
  }

  HepLorentzVector pi0FourMomentum(const Pi0Candidate& pi0, double vtxz) {
    return pi0.fourMomentum(vtxz);
  }

  void FillCesInfo(const Pi0Candidate& pi0, CesInfo& strip1, CesInfo& strip2, 
                   CesInfo& wire1, CesInfo& wire2) {
    const CesMatch* match = pi0.cesMatch();
    CesMatch::MatchType m = match->matchType();

    // initialize stuff
    strip1.rawEnergy = -999.; strip2.rawEnergy = -999.; wire1.rawEnergy = -999.; wire2.rawEnergy = -999.;
    strip1.fittedPosition = -999.; strip2.fittedPosition = -999.; wire1.fittedPosition = -999.; wire2.fittedPosition = -999.;
    strip1.width = -999.; strip2.width = -999.; wire1.width = -999.; wire2.width = -999.;
    strip1.chi2 = -999.; strip2.chi2 = -999.; wire1.chi2 = -999.; wire2.chi2 = -999.;
    strip1.rawPosition = -999.; strip2.rawPosition = -999.; wire1.rawPosition = -999.; wire2.rawPosition = -999.;
    strip1.globalPosition = -999.; strip2.globalPosition = -999.; 
    wire1.globalPosition = -999.; wire2.globalPosition = -999.;
    strip1.fittedEnergy = -999.; strip2.fittedEnergy = -999.;
    wire1.fittedEnergy = -999.; wire2.fittedEnergy = -999.;
    strip1.widthError = -999.; strip2.widthError = -999.;
    wire1.widthError = -999.; wire2.widthError = -999.;
    strip1.distTrackCluster = -999.; strip2.distTrackCluster = -999.;
    wire1.distTrackCluster = -999.; wire2.distTrackCluster = -9999.;
    strip1.halfWedge = -999; strip2.halfWedge = -999; wire1.halfWedge = -999; wire2.halfWedge = -999;
    strip1.firstElement = -999; strip2.firstElement = -999; wire1.firstElement = -999; wire2.firstElement = -999;
    strip1.numElements = -999; strip2.numElements = -999; wire1.numElements = -999; wire2.numElements = -999;

    // mimic Pasha: I don't want to figure out these vectors of vectors of links :-)
    const std::vector< ValueVectorLink <CesClusterColl> >& clist = match->matchedClusters();
    if(m == CesMatch::T1W1S) {
      ValueVectorLink<CesClusterColl> wc1 = clist[0];
      ValueVectorLink<CesClusterColl> sc1 = clist[1];
      strip1.rawEnergy = sc1->raw_energy();
      wire1.rawEnergy = wc1->raw_energy();
      strip1.fittedPosition = sc1->fitted_position();
      wire1.fittedPosition = wc1->fitted_position();
      strip1.width = sc1->width();
      wire1.width = wc1->width();
      strip1.chi2 = sc1->chi2();
      wire1.chi2 = wc1->chi2();
      strip1.rawPosition = sc1->raw_position();
      strip1.globalPosition = sc1->global_position();
      strip1.fittedEnergy = sc1->fitted_energy();
      strip1.widthError = sc1->width_error();
      strip1.distTrackCluster = sc1->dist_track_cluster();
      strip1.halfWedge = sc1->half_wedge();
      strip1.firstElement = sc1->first_element();
      strip1.numElements = sc1->number_elements();
      wire1.rawPosition = wc1->raw_position();
      wire1.globalPosition = wc1->global_position();
      wire1.fittedEnergy = wc1->fitted_energy();
      wire1.widthError = wc1->width_error();
      wire1.distTrackCluster = wc1->dist_track_cluster();
      wire1.halfWedge = wc1->half_wedge();
      wire1.firstElement = wc1->first_element();
      wire1.numElements = wc1->number_elements();
    }
    else if (m == CesMatch::T1W2S_UNSPLIT) {
      ValueVectorLink<CesClusterColl> wc1 = clist[0];
      ValueVectorLink<CesClusterColl> sc1 = clist[1];
      ValueVectorLink<CesClusterColl> sc2 = clist[2];
      strip1.rawEnergy = sc1->raw_energy();
      strip2.rawEnergy = sc2->raw_energy();
      wire1.rawEnergy = wc1->raw_energy();
      strip1.fittedPosition = sc1->fitted_position();
      strip2.fittedPosition = sc2->fitted_position();
      wire1.fittedPosition = wc1->fitted_position();
      strip1.width = sc1->width();
      strip2.width = sc2->width();
      wire1.width = wc1->width();
      strip1.chi2 = sc1->chi2();
      strip2.chi2 = sc2->chi2();
      wire1.chi2 = wc1->chi2();
      strip1.rawPosition = sc1->raw_position();
      strip1.globalPosition = sc1->global_position();
      strip1.fittedEnergy = sc1->fitted_energy();
      strip1.widthError = sc1->width_error();
      strip1.distTrackCluster = sc1->dist_track_cluster();
      strip1.halfWedge = sc1->half_wedge();
      strip1.firstElement = sc1->first_element();
      strip1.numElements = sc1->number_elements();
      strip2.rawPosition = sc2->raw_position();
      strip2.globalPosition = sc2->global_position();
      strip2.fittedEnergy = sc2->fitted_energy();
      strip2.widthError = sc2->width_error();
      strip2.distTrackCluster = sc2->dist_track_cluster();
      strip2.halfWedge = sc2->half_wedge();
      strip2.firstElement = sc2->first_element();
      strip2.numElements = sc2->number_elements();
      wire1.rawPosition = wc1->raw_position();
      wire1.globalPosition = wc1->global_position();
      wire1.fittedEnergy = wc1->fitted_energy();
      wire1.widthError = wc1->width_error();
      wire1.distTrackCluster = wc1->dist_track_cluster();
      wire1.halfWedge = wc1->half_wedge();
      wire1.firstElement = wc1->first_element();
      wire1.numElements = wc1->number_elements();
    }
    else if (m == CesMatch::T2W1S_UNSPLIT) {
      ValueVectorLink<CesClusterColl> sc1 = clist[0];
      ValueVectorLink<CesClusterColl> wc1 = clist[1];
      ValueVectorLink<CesClusterColl> wc2 = clist[2];
      strip1.rawEnergy = sc1->raw_energy();
      wire1.rawEnergy = wc1->raw_energy();
      wire2.rawEnergy = wc2->raw_energy();
      strip1.fittedPosition = sc1->fitted_position();
      wire1.fittedPosition = wc1->fitted_position();
      wire2.fittedPosition = wc2->fitted_position();
      strip1.width = sc1->width();
      wire1.width = wc1->width();
      wire2.width = wc2->width();
      strip1.chi2 = sc1->chi2();
      wire1.chi2 = wc1->chi2();
      wire2.chi2 = wc2->chi2();
      strip1.rawPosition = sc1->raw_position();
      strip1.globalPosition = sc1->global_position();
      strip1.fittedEnergy = sc1->fitted_energy();
      strip1.widthError = sc1->width_error();
      strip1.distTrackCluster = sc1->dist_track_cluster();
      strip1.halfWedge = sc1->half_wedge();
      strip1.firstElement = sc1->first_element();
      strip1.numElements = sc1->number_elements();
      wire1.rawPosition = wc1->raw_position();
      wire1.globalPosition = wc1->global_position();
      wire1.fittedEnergy = wc1->fitted_energy();
      wire1.widthError = wc1->width_error();
      wire1.distTrackCluster = wc1->dist_track_cluster();
      wire1.halfWedge = wc1->half_wedge();
      wire1.firstElement = wc1->first_element();
      wire1.numElements = wc1->number_elements();
      wire2.rawPosition = wc2->raw_position();
      wire2.globalPosition = wc2->global_position();
      wire2.fittedEnergy = wc2->fitted_energy();
      wire2.widthError = wc2->width_error();
      wire2.distTrackCluster = wc2->dist_track_cluster();
      wire2.halfWedge = wc2->half_wedge();
      wire2.firstElement = wc2->first_element();
      wire2.numElements = wc2->number_elements();
    }
    else if (m == CesMatch::T0W1S) {
      ValueVectorLink<CesClusterColl> sc1 = clist[0];
      strip1.rawEnergy = sc1->raw_energy();
      strip1.fittedPosition = sc1->fitted_position();
      strip1.width = sc1->width();
      strip1.chi2 = sc1->chi2();
      strip1.rawPosition = sc1->raw_position();
      strip1.globalPosition = sc1->global_position();
      strip1.fittedEnergy = sc1->fitted_energy();
      strip1.widthError = sc1->width_error();
      strip1.distTrackCluster = sc1->dist_track_cluster();
      strip1.halfWedge = sc1->half_wedge();
      strip1.firstElement = sc1->first_element();
      strip1.numElements = sc1->number_elements();
    }
    else if (m == CesMatch::T1W0S) {
      ValueVectorLink<CesClusterColl> wc1 = clist[0];
      wire1.rawEnergy = wc1->raw_energy();
      wire1.fittedPosition = wc1->fitted_position();
      wire1.width = wc1->width();
      wire1.chi2 = wc1->chi2();
      wire1.rawPosition = wc1->raw_position();
      wire1.globalPosition = wc1->global_position();
      wire1.fittedEnergy = wc1->fitted_energy();
      wire1.widthError = wc1->width_error();
      wire1.distTrackCluster = wc1->dist_track_cluster();
      wire1.halfWedge = wc1->half_wedge();
      wire1.firstElement = wc1->first_element();
      wire1.numElements = wc1->number_elements();
    }
    else if (m == CesMatch::T1W2S_SPLIT) {
    }
    else if (m == CesMatch::T2W1S_SPLIT) {
    }

    return;
  }

  double tauCone(const CdfTau& tau) {
    // What the tau group seems to be doing now:
    if(tau.clusterEt() == 0) {
      return 0.1745; // 10 degrees; website says 0.2, but code says 10deg.
    }
    double r = 5 / tau.clusterEt();
    if(r > 0.1745) r = 0.1745;
    //if(r < .005) r = .005; (Listed on website, but not in code)
    return r;
  }

  bool tauIsTrackInCone(const CdfTau& tau, const CdfTrack& track) {
    Hep3Vector momentum = track.momentum();
    int charge = (track.curvature() > 0) ? 1 : -1;
    double z0 = track.z0();
    double x0 = -1.0 * track.d0() * sin(track.phi0()) * charge;
    double y0 = track.d0() * cos(track.phi0()) * charge;
    HepPoint3D position(x0,y0,z0);
    double trackEta, trackPhi;
    JetVariables::extrapolatedEndEtaPhi(position,momentum,charge,trackEta,trackPhi);

    Hep3Vector smomentum = tau.seedTrack()->momentum();
    int scharge = (tau.seedTrack()->curvature() > 0) ? 1 : -1;
    double sz0 = tau.seedTrack()->z0();
    double sx0 = -1.0 * tau.seedTrack()->d0() * sin(tau.seedTrack()->phi0()) * scharge;
    double sy0 = tau.seedTrack()->d0() * cos(tau.seedTrack()->phi0()) * scharge;
    HepPoint3D sposition(sx0,sy0,sz0);
    double taueta, tauphi;
    JetVariables::extrapolatedEndEtaPhi(sposition,smomentum,scharge,taueta,tauphi);    

    double rsq = (taueta-trackEta)*(taueta-trackEta)+(tauphi-trackPhi)*(tauphi-trackPhi);
    if(rsq < tauCone(tau)*tauCone(tau)) {
      return true;
    }
    return false;
  }

  bool tauIsTrackInConeUnextrapolated(const CdfTau& tau, const CdfTrack& track) {
    double trackEta, trackPhi;
    trackEta = track.pseudoRapidity();
    trackPhi = track.phi0();

    double taueta, tauphi;
    taueta = tau.seedTrack()->pseudoRapidity();
    tauphi = tau.seedTrack()->phi0();

    double rsq = (taueta-trackEta)*(taueta-trackEta)+(tauphi-trackPhi)*(tauphi-trackPhi);
    if(rsq < tauCone(tau)*tauCone(tau)) {
      return true;
    }
    return false;
  }

  bool tauTrackVertexIn5Cm(const CdfTau& tau, const CdfTrack& track) {
    double dist = fabs(tau.seedTrack()->z0() - track.z0());
    if(dist < 5) return true;
    return false;
  }

  // Useful for sorting
  struct TrackDescendingPt {
    bool operator()(const CdfTrack_clnk t1, const CdfTrack_clnk t2)
    {
      return TrackingVariables::trkPt(&*t1) > TrackingVariables::trkPt(&*t2);
    }
  };

  bool tauGetTracks(const CdfTau& tau, CdfTrackView_h& hTrkView, string process)
  {
    // if hTrkView non-empty, clear it, otherwise initialize it
    if (hTrkView.is_null()) {
      hTrkView = new CdfTrackView;
    } else {
      hTrkView->contents().clear();
    }
    
    const CdfTrackView* trackHandle = tau.associatedTauTracks();

    if(trackHandle->contents().begin() == trackHandle->contents().end()) {
      return true; // no error, but an empty track collection. pretty weird.
    }

    // Loop over the tracks
    for(CdfTrackView::const_iterator trackIter = trackHandle->contents().begin();
        trackIter != trackHandle->contents().end(); ++trackIter)
    {
      const CdfTrack& track = **trackIter;
      if(tauIsTrackInCone(tau, track) && tauTrackVertexIn5Cm(tau,track)) {
        hTrkView->contents().push_back( CdfTrack_clnk(&track) );
      }
    }

    std::sort(hTrkView->contents().begin(), hTrkView->contents().end(), TrackDescendingPt());

    return true;
  }

  bool tauGetTracksUnextrapolated(const CdfTau& tau, CdfTrackView_h& hTrkView, string process)
  {
    // if hTrkView non-empty, clear it, otherwise initialize it
    if (hTrkView.is_null()) {
      hTrkView = new CdfTrackView;
    } else {
      hTrkView->contents().clear();
    }
    
    const CdfTrackView* trackHandle = tau.associatedTauTracks();

    if(trackHandle->contents().begin() == trackHandle->contents().end()) {
      return true; // No error, but an empty track collection. Pretty weird.
    }

    // Loop over the tracks
    for(CdfTrackView::const_iterator trackIter = trackHandle->contents().begin();
        trackIter != trackHandle->contents().end(); ++trackIter)
    {
      const CdfTrack& track = **trackIter;
      if(tauIsTrackInConeUnextrapolated(tau, track) && tauTrackVertexIn5Cm(tau,track)) {
        hTrkView->contents().push_back( CdfTrack_clnk(&track) );
      }
    }

    std::sort(hTrkView->contents().begin(), hTrkView->contents().end(), TrackDescendingPt());

    return true;
  }

  int tauTracksInCone(const CdfTau& tau, const CdfTrackView_h thetracks) {
    return thetracks->contents().size();
  }

  int tauTracksChargeInCone(const CdfTau& tau, const CdfTrackView_h thetracks) {
    int charge=0;
    if(thetracks->contents().size() > 0) {
      for(CdfTrackView::const_iterator trkIter = thetracks->contents().begin();
          trkIter != thetracks->contents().end(); ++trkIter) {
        const CdfTrack& track = **trkIter;
        charge += track.charge();
      }
    }
    return charge;
  }

  int tauIsolationTracks(const CdfTau& tau, const CdfTrackView_h thetracks) {
    return (tau.associatedTauTracks()->contents().size() - thetracks->contents().size());
  }

  int tauIsolationPi0s(const CdfTau& tau) {
    return tau.isolationPi0s();
  }

  double tauEHadOverP(const CdfTau& tau) {
    double Ehad = tau.clusterHadronE();
    double P = tau.tracksP();
    if(P == 0) {
      return -9999.0;
    }
    return (Ehad / P);
  }

  double tauCaloMass(const CdfTau& tau) {
    return tau.clusterMass();
  }
} // namespace TauVariables