////////////////////////////////////////////////////////////////////
//
//     File: JetVariables2.cc
//     Authors: Matt Reece
//     Description: High level analysis jet variables defined as functions
//                  This file contains vertex-related variables.
//     Created: 3/26/2001
//
//     Revision: 0.1
//
////////////////////////////////////////////////////////////////////

#include <math.h>
#include "Rtypes.h"
#include "HighLevelObjects/JetVariables.hh"
#include "JetObjects/CdfJet.hh"
#include "CLHEP/Vector/LorentzVector.h"
//#include "Jet/makeJTC96jets.hh"
#include "VertexObjects/Vertex.hh"
#include "VertexObjects/VertexColl.hh"
//#include "JetObjects/BasicJet.hh"
#include "TrackingObjects/Tracks/CdfTrack.hh"
#include "TrackingObjects/Storable/CdfTrackView.hh"
#include "BTagObjects/SecVtxTracks.hh"
#include "ErrorLogger_i/gERRLOG.hh"
#include "CLHEP/Vector/ThreeVector.h"
#include "Extrapolator/MExtrapolator.hh"

namespace JetVariables {
  // 10/11/2001 Tracking Revision
  // Fill hTrkView with the default tracks as given by CdfJet
  bool jetTracksDefault(const CdfJet& jet, CdfTrackView_h& hTrkView, double coneSize, string process)
  {
    /*
    Loop over all the tracks in the collection, testing each one as follows:
    (1) Get the jet's centroidEta, centroidPhi (from CdfJet) as phi_jet,
        eta_jet 
    (2) Use the track extrapolation code to extrapolate to r = r_ces,
        where r_ces is stored in CalorGeometry/CalParameters.hh as TLYCES (add a
        small correction to get to the shower location.) Check whether this is
        central:
        (A) If |z| < z_plug, use r = r_ces, eta, phi at this location
        (B) If |z| > z_plug, use eta, phi at z = z_plug (TLZPEM plus
            correction); we re-extrapolate to this value of z.
    (3) If (eta_track - eta_phi)^2 + (phi_track - phi_jet)^2 < 0.4^2,
        count the track as in the jet.
    (The choice of r_ces is somewhat arbitrary, but is a good starting point.)  
    */

    // if hTrkView non-empty, clear it, otherwise initialize it
    if (hTrkView.is_null()) {
      hTrkView = new CdfTrackView;
    } else {
      hTrkView->contents().clear();
    }

    // follow CdfJet::tracksInCone code to get a list of tracks 
    CdfTrackView_h trackHandle;
    if(CdfTrackView::defTracks(trackHandle,process) != CdfTrackView::OK) {
      ERRLOG(ELwarning,"CdfJet: ") << "Error in CdfJet: could not find track collection" << endmsg;
      trackHandle = 0;
      return false;
    }

    // Get the jet's centroid eta, centroid phi
    double jetPhi = jetCentroidPhi(jet);
    double jetEta = jetCentroidEta(jet);

    // Loop over the tracks
    for(CdfTrackView::const_iterator trackIter = trackHandle->contents().begin();
        trackIter != trackHandle->contents().end(); ++trackIter)
    {
      // Extrapolate to the calorimeter
      const CdfTrack* track = &**trackIter;
      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;
      if (momentum.perp() > 0.5) {
	extrapolatedEndEtaPhi(position,momentum,charge,trackEta,trackPhi);
	
	double phiDiff = (trackPhi - jetPhi);
	double etaDiff = (trackEta - jetEta);
	if(phiDiff > PI) phiDiff -= (2*PI);
	if(phiDiff < -PI) phiDiff += (2*PI);
	if(phiDiff > PI) phiDiff -= (2*PI);
	if(phiDiff < -PI) phiDiff += (2*PI);

	// cone radius
	if (((etaDiff*etaDiff)+(phiDiff*phiDiff))<(coneSize * coneSize))
	  {
	    hTrkView->contents().push_back( CdfTrack_clnk(track) );
	  }
      }
    }

    return true;
  }

  // Get jets in the original manner
  bool jetTracksOriginal(const CdfJet& cdfJet, CdfTrackView_h& hTrkView, double coneSize, string process)
  {    
    CdfTrackView::defTracks(hTrkView, process);
    bool statusOK = cdfJet.tracksInCone(*hTrkView, coneSize);
    if(!statusOK) {
      //ERRLOG(ELwarning, "HLOJet") << "jetTracksOriginal: CdfJet reports no tracks" << endmsg;
      return false;
    }
    
    return true;
  }

  bool extrapolatedEndEtaPhi(HepPoint3D position, Hep3Vector momentum,
                             int charge, double& eta, double& phi)
  {
    // Get CES radius and plug z
    double rCES = TLYCES + 0.20;
    double zPlug = TLZPEM;
 
    // Is it OK to use MExtrapolator here?? It doesn't do weird muon-specific
    // things?
    MExtrapolator extr(CdfDetector::instance());
    extr.setInitialMomentum(momentum);
    extr.setInitialPosition(position);
    extr.setParticleCharge(charge);

    // Extrapolate to the CES radius
    extr.extrapolateToRadius(rCES);
    HepPoint3D loc = extr.getCurrentPosition();

    // Check if we're in the plug, and if so, re-extrapolate there
    if(fabs(loc.z()) > zPlug)
    {
      double zP = (loc.z() > 0 ? 1.0 : -1.0) * (zPlug + 0.20);
   
      extr.setInitialMomentum(momentum);
      extr.setInitialPosition(position);
      extr.extrapolateToZ(zP);

      loc = extr.getCurrentPosition();
    }

    // Now, calc (eta, phi) of loc and store them in the variables we were passed
    double x = loc.x(), y = loc.y(), z = loc.z();
    double r = sqrt(x*x+y*y);
    phi = atan2(y,x);
    double theta = atan2(r,z);
    eta = -log(tan(theta/2.0));

    return true;
  }

  // Revised 3/19/2001 Return both z vertex and sigma in one function
  bool jetZVertexInfo(const CdfJet& jet, const CdfTrackView_h& trks,
                      double& zVertex, double& zSigma)
  {

    // default return values
    zVertex = -999999.0;
    zSigma =  999.0;

    if(trks->contents().begin() == trks->contents().end())
    {
      //ERRLOG(ELwarning, "HLOJet") << "jetZVertexInfo: no tracks" << endmsg;
      return false;
    }  
    // Compute the weighted mean
    double weightedSumZ0 = 0.0;
    double weightSum = 0.0;
    for(CdfTrackView::const_iterator itrack = trks->contents().begin();
        itrack != trks->contents().end(); ++itrack)
    {
      // Is there any way to find the mean and sigma in the same loop?
      const CdfTrack& track = **itrack;

      if(track.getHelixFit().getErrorMatrix()
	 [SimpleReconstructedTrack::Z0][SimpleReconstructedTrack::Z0] <=0.0) {
        ERRLOG(ELerror, "HLOJet") 
	  << "Found CdfTrack with sigma**2 < 0  Track ID " 
	  << track.id() << endmsg;
	return false;
      }

      double sigma = track.sigmaZ0();      

      // Avoid divide by zero errors
      double sigmaSq = sigma * sigma;
      if(sigmaSq == 0.0) {
        ERRLOG(ELerror, "HLOJet") << "Unable to compute z vertex, sigma=0." << endmsg;
        return false;
      }

      double weight = 1.0 / (sigma * sigma);
      weightedSumZ0 += weight * track.z0();
      weightSum += weight;
    }

    // Avoid divide by zero error; should never happen since 1/(sigma^2) > 0,
    // and we already checked for the zero tracks case.
    if(weightSum == 0.0) {
      ERRLOG(ELerror, "HLOJet") << "Unable to compute z vertex, sum of weights is 0." << endmsg;
      return false;
    }

    zVertex = ( weightedSumZ0 / weightSum );

    // Find standard deviation
    int nChargedTracks = trks->contents().size();
    if(nChargedTracks == 1) {
      // One track shouldn't have variance
      //ERRLOG(ELwarning, "HLOJet") << "jetZVertexInfo: only one track, reporting sigma=0" << endmsg;
      zSigma = 0.0;
      return true;
    }

    // Compute the variance
    double sumZ0SquareDiff = 0.0;
    double trackZDiff;
    for(CdfTrackView::const_iterator itrack = trks->contents().begin(); 
        itrack != trks->contents().end(); ++itrack)
    {
      const CdfTrack& track = **itrack;
      trackZDiff = track.z0() - zVertex;
      sumZ0SquareDiff += trackZDiff * trackZDiff;
    }
    double variance = sumZ0SquareDiff / (nChargedTracks - 1);

    // Get standard deviation
    zSigma = sqrt(variance);

    return true;
  }

  // JTZVTX(NJT) Z vertex position from tracks pointing at jet
  // Computes mean from all tracks
  // Modified  1/19/2001: 
  // Weights by 1/sigma^2, not just a simple mean
  double jetZVertex(const CdfJet& jet, const CdfTrackView_h& trks){
    if(trks->contents().begin() == trks->contents().end())
    {
      //ERRLOG(ELwarning, "HLOJet") << "jetZVertex: no tracks" << endmsg;
      return -99999.0;
    }  

    double z, s;
    bool bProblem = !jetZVertexInfo(jet, trks, z, s);
    if(bProblem)
    {
      // scream
      //ERRLOG(ELwarning, "HLOJet") << "Could not find z vertex and sigma for jet";
      return -99999.0;
    }
    return z;
  } 

  // JTDZVTX(NJT)  RMS of Z0's from tracks pointing at jet
  // returns standard deviation
  double jetRmsZ0(const CdfJet& jet, const CdfTrackView_h& trks){
    if(trks->contents().begin() == trks->contents().end())
    {
      //ERRLOG(ELwarning, "HLOJet") << "jetRmsZ0: no tracks" << endmsg;
      return -99999.0;
    }  

    double z, s;
    bool bProblem = !jetZVertexInfo(jet, trks, z, s);
    if(bProblem)
    {
      // scream
      //ERRLOG(ELwarning, "HLOJet") << "Could not find z vertex and sigma for jet";
      return -99999.0;
    }
    return s;
  }

  // 2/1/2001
  // Made VertexColl& a parameter to the functions
  // jetVPt and jetNoVPt. The list of vertex
  // candidates is to be obtained from VxPrim.

  // JTVPT is the sum of the pt of the 3D tracks which
  // point at both the vertex candidate and jet candidate.
  // The VxPrim module should have been called to get
  // the VertexColl parameter.
  double jetVPt(const CdfJet& jet, const CdfTrackView_h& trks, 
		const Vertex& vertex, const VertexColl& allVertices) {
    double sumPt, fractionOccupancy;
    int occupancy;
    bool bResult = jetVtxStats(jet, trks, vertex, allVertices, sumPt, 
			       occupancy, fractionOccupancy);
    if(bResult) {
      return sumPt;
    }
    return -9999.;
  }


  // JTNOVPT is the sum of the pt of the 3D tracks
  // which point at the jet candidate but do not lie 
  // within 5 cm of any vertex candidate.
  // VxPrim should have been used to find vertex
  // candidates.
  double jetNoVPt(const CdfJet& jet,  const CdfTrackView_h& trks, 
                  const VertexColl& allVertices)
  {
    if(trks->contents().begin() == trks->contents().end())
    {
      //ERRLOG(ELwarning, "HLOJet") << "jetNoVPt: no tracks" << endmsg;
      return -99999.0;
    }

    double sumPt = 0.0;  

    // Loop through tracks, adding Pt from the ones which point at no vertex
    for(CdfTrackView::const_iterator itrack = trks->contents().begin();
        itrack != trks->contents().end(); ++itrack)
    {
      const CdfTrack& track = **itrack;
      // Check to see if the track points to any vertex
      bool vertexWithin5Cm = false;
      double trackZ0 = track.z0();
      for(VertexColl::const_iterator ivtx = allVertices.contents().begin();
          ivtx != allVertices.contents().end(); ++ivtx)
      {
        const Vertex& vtx = **ivtx;
        if(vtx.type() == Vertex::Primary)
        {
          double diff = fabs(vtx.z() - trackZ0);
          if(diff < 5) // A vertex is nearby
          {
            vertexWithin5Cm = true;
            break;
          }
        }
      }
      if(vertexWithin5Cm == false) sumPt += track.pt();
    }
    return sumPt;
  }

  // 2/14/2001
  // JTNJTRK0: Number of tracks in the jet used by SECVTX
  int jetNumSecVtxTracks(const CdfJet& jet, const CdfTrackView_h& trks, const char *theDescription)
  {
    if(trks->contents().begin() == trks->contents().end())
    {
      //ERRLOG(ELwarning, "HLOJet") << "jetNumSecVtxTracks: no tracks" << endmsg;
      return -99999;
    }  

    // Get a list of SecVtx tracks (computed by SecVtxModule)
    SecVtxTracks_ch theHandle;
    //std::string description = "SecVtxTracks";
    SecVtxTracks::Error theStatus = SecVtxTracks::find(theHandle, theDescription);
    if(theStatus == SecVtxTracks::ERROR)
    {
      ERRLOG(ELwarning, "HLOJet") << "Error in jetNumSecVtxTracks: can't get SecVtxTracks" << endmsg;
      return -99999;
    }

    // Count the ones that match
    int numSecVtxTracks = theHandle->getntracks();
    int numMatchingTracks = 0;

    // Loop over tracks in the jet
    for(CdfTrackView::const_iterator itrack = trks->contents().begin();
        itrack != trks->contents().end(); ++itrack)
    {
      // Find the highest-level parent of this track.
      const CdfTrack & track = **itrack;
      CdfTrack_clnk parentLnk = track.parent();
      const CdfTrack * parentPtr = parentLnk.operator->();
      const CdfTrack * currentPtr = &track;
      while(parentPtr!=0)
      {
        currentPtr = parentPtr;
        parentLnk =currentPtr->parent();
        parentPtr = parentLnk.operator->();
      }     
      bool matches = false;

      // Loop over all the secondary vertex tracks until we find a match
      for(int i = 0; (i<numSecVtxTracks) && !matches; ++i)
      {
        CdfTrack_clnk iSvtxTrack = theHandle->getrack(i);

        // Compare IDs of the parent tracks of the SvtxTrack and the jet track
        CdfTrack_clnk svxParentLnk = iSvtxTrack->parent();
        const CdfTrack * svxParentPtr = svxParentLnk.operator->();
        const CdfTrack * svxCurrentPtr = iSvtxTrack.operator->();
        while(svxParentPtr != NULL)
        {
          svxCurrentPtr = svxParentPtr;
          svxParentLnk = svxCurrentPtr->parent();
          svxParentPtr = svxParentLnk.operator->();
        }

        if(svxCurrentPtr->id() == currentPtr->id())
        {
          ++numMatchingTracks;
          matches = true; // this stops the loop and goes to next jet track
        }
      }
    }

    return numMatchingTracks;
  }

} // namespace JetVariables