////////////////////////////////////////////////////////////////////
//   File:  TauVariables2.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 "TauObjects/Pi0CandidateView.hh"
#include "TauObjects/TauCluster.hh"
#include "CalorObjects/PhysicsTower.hh"
#include "AbsEnv/AbsEnv.hh"
#include <math.h>
#include <algorithm>

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

  struct Pi0DescendingPtSorter {
    double refvtxz;

    bool operator()(const Pi0CandidateView::Pi0Candidate_link pi1, 
                    const Pi0CandidateView::Pi0Candidate_link pi2)
    {
      HepLorentzVector v1 = pi0FourMomentum(*pi1, refvtxz);
      HepLorentzVector v2 = pi0FourMomentum(*pi2, refvtxz);
      double pt1sq = v1.px()*v1.px()+v1.py()*v1.py();
      double pt2sq = v2.px()*v2.px()+v2.py()*v2.py();
      return pt1sq > pt2sq;
    }
  };

  struct TrackAscendingNearnessSorter {
    Hep3Vector seed;

    bool operator()(const CdfTrack_clnk t1, const CdfTrack_clnk t2)
    {
      double angle1 = seed.angle(t1->momentum());
      double angle2 = seed.angle(t2->momentum());
      return angle1 < angle2;
    }
  };

  struct Pi0AscendingNearnessSorter {
    Hep3Vector seed;
    double refvtxz;

    bool operator()(const Pi0CandidateView::Pi0Candidate_link pi1,
                    const Pi0CandidateView::Pi0Candidate_link pi2)
    {
      HepLorentzVector v1 = pi0FourMomentum(*pi1, refvtxz);
      HepLorentzVector v2 = pi0FourMomentum(*pi2, refvtxz);
      return seed.angle(v1) < seed.angle(v2);
    }
  };

  bool tauTrackGetter(const CdfTau& tau, CdfTrackView_h& tracks,
	double inner, double outer, double threshold,
                      bool nearsort, double& nearangle)
  {
    if(tracks.is_null()) {
      tracks = new CdfTrackView;
    } else {
      tracks->contents().clear();
    }

    // GET ALL THE TRACKS FROM DEFTRACKS: mimic TauCandidate.cc
    CdfTrackView_h handle;
    AbsEvent* event = AbsEnv::theEvent();
    bool OK = 
      CdfTrackView::defTracks (handle, event->process_name ()) 
        == CdfTrackView::OK ||
      CdfTrackView::defTracks (handle, "PROD") == CdfTrackView::OK ||
      CdfTrackView::defTracks (handle, "L3") == CdfTrackView::OK;
    if(!OK) return false;

    Hep3Vector seed = tau.seedTrack()->momentum();
    for(CdfTrackView::const_iterator iter = handle->contents().begin();
	iter != handle->contents().end(); ++iter)
    {
      const CdfTrack& track = **iter;
      double angle = seed.angle(track.momentum());

      // we want to use the tau group's vertex criterion
      double refz = tau.referenceVertexZ();
      double trackz = track.z0();

      if(inner <=  angle && angle < outer && fabs(trackz-refz)<10.
	 && track.pt() >= threshold && track.numCTHitsAx() >= 10
         && track.numCTHitsSt() >= 10)
      {
        tracks->contents().push_back( CdfTrack_clnk(&track) );
      }
    }

    if(tracks->contents().size() > 0) {
      if(!nearsort) {
        std::sort(tracks->contents().begin(), tracks->contents().end(),
                  TrackDescendingPtSorter());
      }
      else {
        // sort by proximity to the seed
        TrackAscendingNearnessSorter sorter;
        sorter.seed = tau.seedTrack()->momentum();
        std::sort(tracks->contents().begin(), 
                  tracks->contents().end(),
                  sorter);
        // So ugly! The mind boggles. (Do other things boggle?)
        nearangle = sorter.seed.angle((*tracks->contents().begin())->momentum());
      }
    }

    return true; 
  }	

  bool tauGetConeTracks(const CdfTau& tau, CdfTrackView_h& theconetracks,
 	double conesize, double trkconethreshold)
  {
    double dummy;
    return tauTrackGetter(tau, theconetracks, 0., conesize, 
	trkconethreshold, false, dummy);	
  }

  bool tauGetIsoTracks(const CdfTau& tau, CdfTrackView_h& theisotracks,
                       double conesize, double annulus, 
                       double trkisothreshold)
  {
    double dummy;
    return tauTrackGetter(tau, theisotracks, conesize, annulus,
	trkisothreshold, false, dummy);	
  }

  bool tauGetNearTracks(const CdfTau& tau, CdfTrackView_h& theneartracks,
 	double conesize, double trkisothreshold, double& nearangle)
  {
    return tauTrackGetter(tau, theneartracks, conesize, 100000., 
	trkisothreshold, true, nearangle);	
  }

  bool tauPi0Getter(const CdfTau& tau, Pi0CandidateView& thepi0s,
		    double inner, double outer, double threshold,
                    bool nearsort, double& nearangle)
  {
    // Mimic the code in TauCandidate::matchPi0s
    Pi0CandidateColl_ch pi0s;
    AbsEvent* event = AbsEnv::theEvent();
    bool OK = Pi0CandidateCollection::find(&pi0s, *event, "");
    if(!OK) return false;

    Hep3Vector seed = tau.seedTrack()->momentum();
    double etaS = -log(tan(seed.theta()/2.0));
    double phiS = seed.phi();
    Pi0CandidateCollection::const_iterator pi0 = pi0s->contents().begin();
    for( ; pi0 != pi0s->contents().end(); pi0++) {
      HepLorentzVector pi0_p4;
      if (pi0->fillFourMomentum(&pi0_p4, tauReferenceVtxZ(tau))) {
        if(pi0_p4.et() < threshold) continue;

        if(inner <= pi0_p4.angle(seed) && pi0_p4.angle(seed) < outer) {
          thepi0s.add(Pi0CandidateView::Pi0Candidate_link(pi0s,pi0));
        }
      }
    }

    if(thepi0s.contents().size() > 0) {
      if(!nearsort) {
        Pi0DescendingPtSorter sorter;
        sorter.refvtxz = tauReferenceVtxZ(tau);
        std::sort(thepi0s.contents().begin(), thepi0s.contents().end(),
                  sorter);
      }
      else {
        Pi0AscendingNearnessSorter sorter;
        sorter.refvtxz = tauReferenceVtxZ(tau);
        sorter.seed = tau.seedTrack()->momentum();
        std::sort(thepi0s.contents().begin(), thepi0s.contents().end(),
                  sorter);
        
        HepLorentzVector nearmomentum = 
          pi0FourMomentum(**thepi0s.contents().begin(), sorter.refvtxz);
        nearangle = sorter.seed.angle(nearmomentum);
      }
    }
      
    return true;
  }

  bool tauGetConePi0s(const CdfTau& tau, Pi0CandidateView& theconepi0s,
		      double conesize, double pi0conethreshold)
  {
    double dummy;
    return tauPi0Getter(tau, theconepi0s, 0., conesize, pi0conethreshold,
                        false, dummy);
  }

  bool tauGetIsoPi0s(const CdfTau& tau, Pi0CandidateView& theisopi0s,
		     double conesize, double annulus, double pi0isothreshold)
  {
    double dummy;
    return tauPi0Getter(tau, theisopi0s, conesize, annulus, pi0isothreshold,
                        false, dummy);
  }

  bool tauGetNearPi0s(const CdfTau& tau, Pi0CandidateView& thenearpi0s,
		      double conesize, double pi0isothreshold,
                      double& nearangle)
  {
    return tauPi0Getter(tau, thenearpi0s, conesize, 10000., pi0isothreshold,
                        true, nearangle);
  }

  double tauMaxD0(const CdfTau& tau, CdfTrackView_h theconetracks)
  {
    double d0 = 0.;
    const CdfTrackView::collection_type& tracks = theconetracks->contents();
    if(tracks.size() > 0) {
      for(CdfTrackView::const_iterator conetrackiter = tracks.begin();
          conetrackiter != tracks.end(); ++conetrackiter)
      {
        double thisd0 = (*conetrackiter)->d0();
        if(fabs(thisd0) > fabs(d0)) d0 = thisd0;
      }
    }
    return d0;
  }

  double tauScalarPtSum(const CdfTau& tau, CdfTrackView_h theconetracks)
  {
    double pt = 0.;
    const CdfTrackView::collection_type& tracks = theconetracks->contents();
    if(tracks.size() > 0) {
      for(CdfTrackView::const_iterator conetrackiter = tracks.begin();
          conetrackiter != tracks.end(); ++conetrackiter)
      {
        pt += TrackingVariables::trkPt(&**conetrackiter);
      }
    }
    return pt;
  }

  double tauScalarPSum(const CdfTau& tau, CdfTrackView_h theconetracks)
  {
    double p = 0.;
    const CdfTrackView::collection_type& tracks = theconetracks->contents();
    if(tracks.size() > 0) {
      for(CdfTrackView::const_iterator conetrackiter = tracks.begin();
          conetrackiter != tracks.end(); ++conetrackiter)
      {
        p += TrackingVariables::trkP(&**conetrackiter);
      }
    }
    return p;
  }

  double tauScalarVisPtSum(const CdfTau& tau, CdfTrackView_h theconetracks,
                           Pi0CandidateView& conepi0s)
  {
    double pt = tauScalarPtSum(tau, theconetracks);
    const Pi0CandidateView::collection_type& pi0s = conepi0s.contents();
    if(pi0s.size() > 0) {
      for(Pi0CandidateView::collection_type::const_iterator 
            conepi0iter = pi0s.begin();
          conepi0iter != pi0s.end(); ++conepi0iter)
      {
        pt += pi0FourMomentum(**conepi0iter, tauReferenceVtxZ(tau)).perp();
      }
    }
    return pt;
  }

  double tauScalarVisPSum(const CdfTau& tau, CdfTrackView_h theconetracks,
                           Pi0CandidateView& conepi0s)
  {
    double p = tauScalarPSum(tau, theconetracks);
    const Pi0CandidateView::collection_type& pi0s = conepi0s.contents();
    if(pi0s.size() > 0) {
      for(Pi0CandidateView::collection_type::const_iterator 
            conepi0iter = pi0s.begin();
          conepi0iter != pi0s.end(); ++conepi0iter)
      {
        p += pi0FourMomentum(**conepi0iter, tauReferenceVtxZ(tau)).e();
      }
    }
    return p;
  }

  // sum of 4-momenta from tracks
  HepLorentzVector tauTracksMomentum(const CdfTau& tau, 
                                     CdfTrackView_h conetracks)
  {
    HepLorentzVector total;
    const CdfTrackView::collection_type& tracks = conetracks->contents();
    if(tracks.size() > 0) {
      for(CdfTrackView::const_iterator conetrackiter = tracks.begin();
          conetrackiter != tracks.end(); ++conetrackiter)
      {
        Hep3Vector momentum = TrackingVariables::trkMomentum(&**conetrackiter);
        HepLorentzVector p(momentum, momentum.mag());
        total += p;
      }
    }
    return total;
  }

  // sum of 4-momenta from tracks and from pi0s
  HepLorentzVector tauVisibleMomentum(const CdfTau& tau,
                                      CdfTrackView_h conetracks,
                                      Pi0CandidateView& conepi0s)
  {
    double refvtxz = tauReferenceVtxZ(tau);
    HepLorentzVector total(0.,0.,0.,0.);
    const CdfTrackView::collection_type& tracks = conetracks->contents();
    if(tracks.size() > 0) {
      for(CdfTrackView::const_iterator conetrackiter = tracks.begin();
          conetrackiter != tracks.end(); ++conetrackiter)
      {
        Hep3Vector momentum = TrackingVariables::trkMomentum(&**conetrackiter);
        HepLorentzVector p(momentum, momentum.mag());
        total += p;
      }
    } 
    const Pi0CandidateView::collection_type& pi0s = conepi0s.contents();
    if(pi0s.size() > 0) {
      for(Pi0CandidateView::collection_type::const_iterator 
            conepi0iter = pi0s.begin();
          conepi0iter != pi0s.end(); ++conepi0iter)
      {
        total += pi0FourMomentum(**conepi0iter, refvtxz);
      }
    }

    return total;
  }

  double tauTowerEt(const CdfTau& tau, int tower)
  {
    return tau.towerEt(tower);
  }

  // mimic the CdfTau::towerEt function
  double tauTowerEmEt(const CdfTau& tau, int tower)
  {
    const TauCluster* cluster = tau.caloCluster();
    if(cluster) {
      int ntowers = cluster->towers().size();
      if(ntowers > 0 && tower < ntowers) {
        // Sort the towers in order of decreasing total Et
        // Can't sort the collection in place, since all
        // this stuff is const...
        std::vector<float> ets(ntowers);
        for(int i = 0; i < ntowers; i++) {
          ets[i] = (cluster->towers().contents()[i])->totEt();
        }
        std::sort(ets.begin(), ets.end());
        for(int j = 0; j < ntowers; j++) {
          double et = (cluster->towers().contents()[j])->totEt();
          if(et == ets[tower]) {
            return (cluster->towers().contents()[j])->emEt();
          }
        }
      }
    }
    return -1.;
  }

  double tauTowerHadEt(const CdfTau& tau, int tower)
  {
    const TauCluster* cluster = tau.caloCluster();
    if(cluster) {
      int ntowers = cluster->towers().size();
      if(ntowers > 0 && tower < ntowers) {
        // Sort the towers in order of decreasing total Et
        // Can't sort the collection in place, since all
        // this stuff is const...
        std::vector<float> ets(ntowers);
        for(int i = 0; i < ntowers; i++) {
          ets[i] = (cluster->towers().contents()[i])->totEt();
        }
        std::sort(ets.begin(), ets.end());
        for(int j = 0; j < ntowers; j++) {
          double et = (cluster->towers().contents()[j])->totEt();
          if(et == ets[tower]) {
            return (cluster->towers().contents()[j])->hadEt();
          }
        }
      }
    }
    return -1.;
  }

  struct TowerDescendingEtSorter {
    bool operator()(const double Et1, const double Et2)
    {
      return Et1 > Et2;
    }
  };

  int tauTowerIEta(const CdfTau& tau, int tower)
  {
    const TauCluster* cluster = tau.caloCluster();
    if(cluster) {
      int ntowers = cluster->towers().size();
      if(ntowers > 0 && tower < ntowers) {
        // Sort the towers in order of decreasing total Et
        // Can't sort the collection in place, since all
        // this stuff is const...
        std::vector<float> ets(ntowers);
        for(int i = 0; i < ntowers; i++) {
          ets[i] = (cluster->towers().contents()[i])->totEt();
        }
        std::sort(ets.begin(), ets.end());
        for(int j = 0; j < ntowers; j++) {
          double et = (cluster->towers().contents()[j])->totEt();
          if(et == ets[tower]) {
            return (cluster->towers().contents()[j])->iEta();
          }
        }
      }
    }
    return -999;
  }

  int tauTowerIPhi(const CdfTau& tau, int tower)
  {
    const TauCluster* cluster = tau.caloCluster();
    if(cluster) {
      int ntowers = cluster->towers().size();
      if(ntowers > 0 && tower < ntowers) {
        // Sort the towers in order of decreasing total Et
        // Can't sort the collection in place, since all
        // this stuff is const...
        std::vector<float> ets(ntowers);
        for(int i = 0; i < ntowers; i++) {
          ets[i] = (cluster->towers().contents()[i])->totEt();
        }
        std::sort(ets.begin(), ets.end());
        for(int j = 0; j < ntowers; j++) {
          double et = (cluster->towers().contents()[j])->totEt();
          if(et == ets[tower]) {
            return (cluster->towers().contents()[j])->iPhi();
          }
        }
      }
    }
    return -999;
  }
} // namespace TauVariables