////////////////////////////////////////////////////////////////////
//
//     File: JetVariables.cc
//     Authors: Abe DeBenedetti, Matt Reece
//     Description: High level analysis jet variables defined as functions
//     Created: 11/15/2000
//
//     Revision: 0.1
//
////////////////////////////////////////////////////////////////////

//#include <cstdio>
#include <math.h>
#include "Rtypes.h"
#include "HighLevelObjects/JetVariables.hh"
#include "JetObjects/CdfJet.hh"
//#include "CLHEP/Geometry/Point3D.h"
#include "CLHEP/Vector/LorentzVector.h"
//#include "Jet/makeJTC96jets.hh"
#include "VertexObjects/Vertex.hh"
#include "VertexObjects/VertexColl.hh"
//#include "JetObjects/BasicJet.hh"
#include "JetUser/JetEnergyCorrections.hh"
#include "TrackingObjects/Tracks/CdfTrack.hh"
#include "TrackingObjects/Storable/CdfTrackView.hh"
#include "BTagObjects/SecVtxTracks.hh"
#include "ErrorLogger_i/gERRLOG.hh"
#include "ElectronObjects/CdfEmObjectView.hh"
//#include "HighLevelObjects/PhotonVariables.hh"
#include "MetObjects/CdfMet.hh"
#include "CalorGeometry/TowerGeometry.hh"
#include "CalorGeometry/Locations.hh"

namespace JetVariables {

  // WARNING! Jet z-vertex corrections here are rough!
  // For best accuracy, recluster.
  //
  // cosine theta, corrected for z vertex
  // (Use cosine to be sure we get the right signs)
  double jetCosTheta(const CdfJet& jet, double zVertex) {
    bool zCorrect = (zVertex > NO_Z_CORRECTION + 1.0);
    if(zCorrect)
    {
      TowerGeometry tGeom;
      double etaDet = jetDetectorEta(jet);
      size_t iEta = tGeom.iEta(etaDet);
      Locations loc(zVertex);
      double emCosTheta = loc.emCosTheta(iEta);
      double hadCosTheta = loc.hadCosTheta(iEta);
      double emFrac = jetEmFraction(jet);
      double weightedCosTheta = emFrac*emCosTheta + (1.0-emFrac)*hadCosTheta;
      return weightedCosTheta;
    }
    return cos(jet.theta());
  }

  // sine theta, corrected for z vertex
  double jetSinTheta(const CdfJet& jet, double zVertex) {
    bool zCorrect = (zVertex > NO_Z_CORRECTION + 1.0);
    if(zCorrect)
    {
      double cos = jetCosTheta(jet, zVertex);
      return sin(acos(cos));
    }
    return sin(jet.theta());
  }

  //what was writen before
  double jetEt(const CdfJet& jet, double zVertex) {
    return jet.totEnergy() * jetSinTheta(jet, zVertex);
  }

  double jetMass(const CdfJet& jet){
    return jet.mass();
  }

  double jetEta(const CdfJet& jet, double zVertex){
    bool zCorrect = zVertex > NO_Z_CORRECTION + 1.0;
    if(zCorrect) {
      double theta = jetTheta(jet, zVertex);
      return -log(tan(theta/2.0));
    }
    return jet.eta();
  }

  double jetDetectorEta(const CdfJet& jet){
    return jet.etaDetector();
  }

  double jetPhi(const CdfJet& jet){
    return jet.phi();
  }

  double jetEmFraction(const CdfJet& jet){
    return jet.emFraction();
  }
  
  int jetNChargedTracks(const CdfJet& jet, const CdfTrackView_h& trks) {    
    return trks->contents().size();
  }

  double jetEp(const CdfJet& jet){
    // Make sure we don't have a divide by zero error
    if(jet.totP() == 0.0)
    {
      ERRLOG(ELerror, "HLOJet") << "jet.totP() = 0, unable to find jetEp." << endmsg;
      return -99999.0;
    }

    return (jet.totEnergy()/jet.totP());
  }

  double jetGuardEnergy(const CdfJet& jet){
    return jet.guardEnergy();
  }

  double jetChargeFraction(const CdfJet& jet, const CdfTrackView_h& trks, 
                           double zVertex) 
  {
    if(trks->contents().begin() == trks->contents().end())
    {
      //ERRLOG(ELwarning, "HLOJet") << "jetChargeFraction: no tracks" << endmsg;
      return -99999.0;
    }  
    double et, sumPt;
    if((et = jetEt(jet,zVertex)) == 0.0)
    {
      ERRLOG(ELerror, "HLOJet") << "jetEt() == 0, unable to find charge fraction." << endmsg;
      return -99999.0;
    }
    if((sumPt = jetPtOut(jet,trks)) <= -99999.0+1)
    {
      //ERRLOG(ELwarning, "HLOJet") << "jetPtOut(jet) can't be found, so cannot find charge fraction." << endmsg;
      return -99999.0;
    }
    return sumPt / et;
  }

  double jetCentroidEta(const CdfJet& jet){
    return jet.centroidEta();
  }

  double jetCentroidPhi(const CdfJet& jet){
    return jet.centroidPhi();
  }

  double jetEtaEtaMoment(const CdfJet& jet){
    return jet.etaEtaMoment();
  }

  double jetEtaPhiMoment(const CdfJet& jet){
    return jet.etaPhiMoment();
  }

  double jetPhiPhiMoment(const CdfJet& jet){
    return jet.phiPhiMoment();
  }

  double jetCentroidEt(const CdfJet& jet){
    return jet.centroidEt();
  }

  double jetTotP(const CdfJet& jet){
    return jet.totP();
  }

  double jetPt(const CdfJet& jet, double zVertex){
    bool zCorrect = zVertex > NO_Z_CORRECTION+1.0;
    if(zCorrect)
    {
      return jet.totP() * jetSinTheta(jet, zVertex);
    }
    return jet.pt();
  }
  
  double jetPtSquared(const CdfJet& jet, double zVertex){
    bool zCorrect = zVertex > NO_Z_CORRECTION + 1.0;
    if(zCorrect)
    {
      HepLorentzVector p = jetFourMomentum(jet, zVertex);
      return p.x()*p.x()+p.y()*p.y();
    }
    return jet.calculatePtSquared();
  }

  double jetTheta(const CdfJet& jet, double zVertex){
    bool zCorrect = zVertex > NO_Z_CORRECTION + 1.0;
    if(zCorrect)
    {
      double theta = acos(jetCosTheta(jet, zVertex));
      return theta;
    }
    return jet.theta();
  }
 
  double jetMassSquared(const CdfJet& jet){
    return jet.massSquared();
  }

  // Rapidity, calculated in the same way CdfJet does it. This is not 
  // the same as eta; it agrees for massless particles.
  double jetRapidity(const CdfJet& jet, double zVertex){
    bool zCorrect = zVertex > NO_Z_CORRECTION + 1.0;
    if(zCorrect)
    {
      HepLorentzVector p = jetFourMomentum(jet,zVertex);
      if(p.e() == p.pz())
      {
        ERRLOG(ELerror, "HLOJet") << "Cannot give jetRapidity: E = Pz" << endmsg;
      }
      return 0.5*log(p.e()+p.pz())/(p.e()-p.pz());
    }    
    return jet.calculateRapidity();
  }

  HepLorentzVector jetFourMomentum(const CdfJet& jet, double zVertex){
    bool zCorrect = zVertex > NO_Z_CORRECTION + 1.0;
    if(zCorrect)
    {
      double cosTheta = jetCosTheta(jet, zVertex);
      double sinTheta = sin(acos(cosTheta));
      double px = jet.totP()*cos(jetPhi(jet))*sinTheta;
      double py = jet.totP()*sin(jetPhi(jet))*sinTheta;
      double pz = jet.totP()*cosTheta;
      return HepLorentzVector(px,py,pz,jetTotEnergy(jet));
    }
    return jet.fourMomentum();
  }

  double jetPx(const CdfJet& jet, double zVertex){
    bool zCorrect = zVertex > NO_Z_CORRECTION + 1.0;
    if(zCorrect)
    {
      return jet.totP()*cos(jetPhi(jet))*jetSinTheta(jet, zVertex);
    }
    return jet.px();
  }

  double jetPy(const CdfJet& jet, double zVertex){
    bool zCorrect = zVertex > NO_Z_CORRECTION + 1.0;
    if(zCorrect)
    {
      return jet.totP()*sin(jetPhi(jet))*jetSinTheta(jet, zVertex);
    }
    return jet.py();
  }

  double jetPz(const CdfJet& jet, double zVertex){
    bool zCorrect = zVertex > NO_Z_CORRECTION + 1.0;
    if(zCorrect)
    {
      return jet.totP()*jetCosTheta(jet,zVertex);
    }
    return jet.pz();
  }

  double jetTotEnergy(const CdfJet& jet){
    return jet.totEnergy();
  }

  // JT_PTOUT(NJT) Sum Pt of tracks pointing at Jet
  double jetPtOut(const CdfJet& jet, const CdfTrackView_h& trks){
    if(trks->contents().begin() == trks->contents().end())
    {
      //ERRLOG(ELwarning, "HLOJet") << "jetPtOut: no tracks" << endmsg;
      return -99999.0;
    }  
    // Add up Pt from each track
    double sumPt = 0.0;
    for(CdfTrackView::const_iterator itrack = trks->contents().begin(); 
        itrack != trks->contents().end(); ++itrack)
    {
      const CdfTrack& track  = **itrack;
      // Access the "pt" member of the track (defined in class
      // ReconstructedTrack, which CdfTrack inherits) 
      sumPt += track.pt();
    }

    return sumPt;
  }

  // JTNTRP(NJT) 
  // Number of tracks in the jet with positive impact parameter
  // 3/28/2001 Modified this; the sign of d0 as reported by the track
  // gives the sign of the angular momentum about the z axis.
  // This is not what we seek. Instead, we test whether the track momentum
  // lies within ninety degrees of the jet momentum; if so, the impact
  // parameter is considered positive. Otherwise, there is a discrepancy
  // between the direction of the track and the direction of the jet
  // to which it is assigned.
  int jetNTracksPositiveIP(const CdfJet& jet, const CdfTrackView_h& trks)
  {
    if(trks->contents().begin() == trks->contents().end())
    {
      //ERRLOG(ELwarning, "HLOJet") << "jetNTracksPositiveIP: no tracks" << endmsg;
      return -99999;
    }  

    // Loop over tracks; count the ones with impact parameter > 0
    int ntrp = 0;
    for(CdfTrackView::const_iterator itrack = trks->contents().begin();
        itrack != trks->contents().end(); ++itrack)
    {
      const CdfTrack & track = **itrack;

      // Get the jet momentum in the x-y plane
      double pXJet = jetPx(jet);
      double pYJet = jetPy(jet);

      // Extract relevant track parameters.
      double d0 = track.d0();
      double phi0 = track.phi0();
      double charge = track.curvature() > 0 ? 1.0 : -1.0;

      // The momentum of the heavy particle which decayed to give us the track
      // is parallel to the impact parameter vector.
      double dX = - d0 * sin(phi0) * charge;
      double dY = d0 * cos(phi0) * charge;

      // Check the sign of the dot product
      double dotProduct = pXJet * dX + pYJet * dY;
      if(dotProduct > 0) ntrp++;
    }
    return ntrp;
  } 

  // Number of tracks in the jet with positive angular momentum about the z-axis
  int jetNTracksPositiveLz(const CdfJet& jet, const CdfTrackView_h& trks)
  {
    if(trks->contents().begin() == trks->contents().end())
    {
      //ERRLOG(ELwarning, "HLOJet") << "jetNTracksPositiveLz: no tracks" << endmsg;
      return -99999;
    }  

    // Count the ones with Lz > 0; the tracking code gives the sign
    // of d0 based on this, so check that.
    int ntrp = 0;
    for(CdfTrackView::const_iterator itrack = trks->contents().begin();
        itrack != trks->contents().end(); ++itrack)
    {
      const CdfTrack & track = **itrack;
      if(track.d0() < 0) ntrp++;
    }
    return ntrp;
  } 

  // 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.

  // UPDATE 7/1/2002: Now will return occupancy (# of tracks pointing to vertex)
  // as well as percentage occupancy.
  bool jetVtxStats(const CdfJet& jet, const CdfTrackView_h& trks, const Vertex&
vertex,
                   const VertexColl& allVertices, double& sumPt, int& occupancy,
                   double& fractionOccupancy) {
    occupancy = 0;
    fractionOccupancy = 0.0;
    sumPt = 0.0;
    int ntrks = 0;
    if(trks->contents().begin() == trks->contents().end()) {
      //ERRLOG(ELwarning, "HLOJet") << "jetVPt: no tracks" << endmsg;
      return false;
    }
    
    // Loop over each track; decide if it points at the vertex.
    for(CdfTrackView::const_iterator itrack = trks->contents().begin();
        itrack != trks->contents().end(); ++itrack, ++ntrks) {
      const CdfTrack& track  = **itrack;
      
      // Check to see if track also points at the vertex
      bool isCorrectVertex = true;
      double trackZ0 = track.z0();
      double vtxDiff = fabs(vertex.z() - track.z0());
      if(vtxDiff > 5) continue; // Too far from vertex

      // Now, check the other vertices to see if any are closer
      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 < vtxDiff) { // We have a better match
            isCorrectVertex = false;
            break;
          }
        }
      }
      
      // If the track points at the vertex, add its Pt
      if(isCorrectVertex) {
        sumPt += track.pt();
        occupancy++;
      }
    }
    fractionOccupancy = double(occupancy) / double(ntrks);
    
    return true;
  }

  // 2/7/2001

  // JTCORFM: Underlying event correction with event-by-event multiple
  // vertex correction, out-of-cone correction, absolute energy scale
  // correction. 

  // Modified 10/12/2001: It appears new parameters have been added to makeJTC96Jet.
  // I'm adding them here too...

  // Modified 15 Mar 2003: Uses new framework.  Thanks to J.F. Arguin for
  // the code.
  // NOTE: THE MEANING OF nRun HAS CHANGED.  It now does run-by-run
  // corrections.
  // Meaning of parameters:
  // offVer = 2 for data, 0 for Monte Carlo
  // nVertex = # of class 12 ZVertices
  // nRun = run number
  // coneSize = jet cone size as a real number (the function will translate
  //            to 0, 1, 2)
  double jetCorFM(const CdfJet& jet, int offVer, int nVertex, int nRun, double coneSize)
  {
    // Make sure we won't have errors
    if(jet.et() <= 0.0 || jet.et() > 800)
    {
      errlog.setSubroutine("JetVariables::jetCorFM");
      ERRLOG(ELerror, "Jet et out of range") << "Jet correction cannot be computed" << endmsg;
      return -99999.0;
    }

    int level = 7; // full corrections
    int syscode = 0;
    int iConeSize = 0;
    if (coneSize > 0.41) {
      iConeSize = 1;
    } else if (coneSize > 0.71) {
      iConeSize = 2;
    }

    JetEnergyCorrections myJetEnergyCorrections=JetEnergyCorrections("JetCorrections", "JetCorrections", level, nVertex, iConeSize, offVer, syscode, nRun);
    float emfr = jet.emFraction();
    HepLorentzVector jetp = jet.fourMomentum();
    double corrFactor = myJetEnergyCorrections.doEnergyCorrections(jetp, emfr, jet.etaDetector());
    
    // Return the correction factor.
    return ( corrFactor ); 
  }

  // JTCORFD: No underlying event correction no out-of-cone correction, 
  // yes absolute energy scale correction.
  // Modified 15 Mar 2003: Uses new framework.  Thanks to J.F. Arguin for
  // the code.
  // NOTE: THE MEANING OF nRun HAS CHANGED.  It now does run-by-run
  // corrections.
  // Meaning of parameters:
  // offVer = 2 for data, 0 for Monte Carlo
  // nVertex = # of class 12 ZVertices
  // nRun = run number
  // coneSize = jet cone size as a real number (the function will translate
  //            to 0, 1, 2)
  double jetCorFD(const CdfJet& jet, int offVer, int nVertex, int nRun, double coneSize)
  {
    // Make sure we won't have errors
    if(jet.et() == 0.0 || jet.et() > 800)
    {
      errlog.setSubroutine("JetVariables::jetCorFD");
      ERRLOG(ELerror, "Jet et out of range") << "Jet correction cannot be computed" << endmsg;
      return -99999.0;
    }

    int level = 5; // no underlying event corrections, no out-of-cone corrections
    int syscode = 0;
    int iConeSize = 0;
    if (coneSize > 0.41) {
      iConeSize = 1;
    } else if (coneSize > 0.71) {
      iConeSize = 2;
    }

    JetEnergyCorrections myJetEnergyCorrections=JetEnergyCorrections("JetCorrections", "JetCorrections", level, nVertex, iConeSize, offVer, syscode, nRun);
    float emfr = jet.emFraction();
    HepLorentzVector jetp = jet.fourMomentum();
    double corrFactor = myJetEnergyCorrections.doEnergyCorrections(jetp, emfr, jet.etaDetector());
    
    // Return the correction factor.
    return ( corrFactor ); 
  }


  // JTWORD: correspondence between jet and other physics objects in the event
  // WARNING: currently only photon matching (bits 7 and 8) and MET matching
  // (bit 0) are implemented
  unsigned short jetWord(const CdfJet& jet, const CdfTrackView_h& trks,
                         double zVertex, double coneSize)
  {

    //using namespace PhotonVariables;

    /*
    bool zCorrect = zVertex > NO_Z_CORRECTION+1.0;

    double zVtx; // z-vertex to pass to status code
    if(zCorrect)
    {
      zVtx = zVertex; // use the parameter given
    }
    else // use the jet vertex
    {
      zVtx = jetZVertex(jet, trks);
    }
    */

    // Refer to stntuple_hbk_fill.F for the original version.
    unsigned short jtword = 0;


    // remove Photon dependence for now
    /*

    CdfEmObjectView_h hEmObjView;
    if(CdfEmObjectView::defEmObjects(hEmObjView) == CdfEmObjectView::ERROR)
    {
      ERRLOG( ELwarning, "HLOJet") << "jetWord: CdfEmObjectView creation failed" << endmsg;
    }
    else // We have EM objects
    {
      int numEmObjects = hEmObjView->size();

      // ELECTRON/JET CORRESPONDENCE
      // Bit 11: jet corresponds to ELES bank with highest ESTAT
      // Bit 10: jet corresponds to ELES bank with 2nd highest ESTAT
      // Bit 9: jet corresponds to ELES bank with 3rd highest ESTAT

      // ESTAT code not implemented yet

      // PHOTON/JET CORRESPONDENCE
      // Bit 8: jet corresponds to ELES bank with highest PSTAT
      // Bit 7: jet corresponds to ELES bank with second highest PSTAT
      // Sort in order of photon status      
      double photEta[2]; // photEta[0] = eta for highest phostat
      double photPhi[2]; // photPhi[0] = phi for highest phostat
      double phoStatMax = -999.0;

      // Find eta, phi for the two Em objects with highest PSTAT
      for(CdfEmObjectView::const_iterator iterEm = hEmObjView->begin();
          iterEm != hEmObjView->end(); ++iterEm)
      {
        const CdfEmObject& emObj = **iterEm;

        // Note: eventually make zvertex a parameter and pass
        // it to phoStat. Until then, do this:
        double phoStat = phoStatus(emObj, zVtx);
        if(phoStat > phoStatMax)
        {
          phoStatMax = phoStat;
          photEta[1] = photEta[0]; // copy second highest
          photPhi[1] = photPhi[0]; 
          photEta[0] = phoDteta(emObj); // get data for this one
          photPhi[0] = phoPhi(emObj);
        }
      }

      // Check to see if these lie within the jet
      int nPhot = min(2, numEmObjects);
      for(int i=0; i < nPhot; i++)
      {
        // Check to see if the photon object has eta, phi within the jet cone
        double etaDiff = fabs(jetDetectorEta(jet) - photEta[i]);
        double phiDiff = fabs(jetPhi(jet) - photPhi[i]);
        if(phiDiff > PI) phiDiff = 2*PI - phiDiff;
        if(sqrt(phiDiff*phiDiff + etaDiff*etaDiff) < coneSize)
        {
          jtword |= 1 << (8-i);
        }
      }      
    } // end working with EM objects

    */

    // MUON TRACK/JET CORRESPONDENCE
    // Bit 6: the QTRK bank track corresponding the muon w/ highest MUSTAT
    //        has detector eta and phi within the cone radius of jet eta/phi
    // Bit 5: the QTRK for the second muon is within cone
    // Bit 4: The QTRK bank for the third muon is within cone

    // TAU/JET CORRESPONDENCE
    // Bit 3: jet corresponds to tau with highest TSTAT
    // Bit 2: jet corresponds to tau with second highest TSTAT
    // Bit 1: jet corresponds to tau with third highest TSTAT

    //  MET/JET CORRESPONDENCE
    // Bit 0: abs(METPHI4-JTPHI1).lt.coneSize
    // Does this have all the corrections that it should? What if there
    // is more than one MET object?
    CdfMet_ch chMet;
    CdfMet::Error metStatus = CdfMet::find(chMet);
    if(metStatus == CdfMet::ERROR) {
      ERRLOG(ELwarning, "HLOJet") << "jetWord: cannot find CdfMet" << endmsg;
    }
    else {
      double metPhi = chMet->phi();
      double metPhiDiff = fabs(metPhi - jetPhi(jet));
      if(metPhiDiff > PI) metPhiDiff = 2*PI - metPhiDiff;
      if(metPhiDiff < coneSize) jtword |= 1;
    }    

    return jtword;
  }

} // namespace JetVariables