// -*- C++ -*-
#include "Rivet/Rivet.hh"
#include "InvariantJetMass.hh"
#include "Rivet/Tools/Logging.hh"
#include "Rivet/Particle.hh"


namespace Rivet {
  
  
  void InvariantJetMass::calc(const FinalState& fs) {
    calc(fs.particles());
  }
  
  void InvariantJetMass::calc(const vector<Particle>& fsparticles) {
    vector<FourMomentum> fourMomenta;
    fourMomenta.reserve(fsparticles.size());
    foreach (const Particle& p, fsparticles) {
      const FourMomentum p4 = p.momentum();
      fourMomenta.push_back(p4);
     }
    _calcInvariantJetMass(fourMomenta);
  }
  
  void InvariantJetMass::calc(const vector<FourMomentum>& fsmomenta) {
    _calcInvariantJetMass(fsmomenta);
  }
   
  /////////////////////////////////////////////////
  
  
  
  inline bool mod2Cmp(const Vector3& a, const Vector3& b) {
    return a.mod2() > b.mod2();
  }
  
  
  /// Do the general case invariantjetmass calbulation
  void _calcM(const vector<FourMomentum>& momenta, double& m) {

    ///create 2 vector onjects, one storing the momentums of
    ///all the final state particles in an event and the other the energies

    vector<Vector3> p;
    p.reserve(momenta.size());
    foreach (const FourMomentum& v, momenta) {
      p.push_back(v.vector3());
    }

    vector<double> energy;
    energy.reserve(momenta.size());
    foreach (const FourMomentum& v, momenta) {
      energy.push_back(v.E());
    }
    
    Vector3 psum;
    
    m = 0.0;
    double en =0.0;
    double pn = 0.0;
    ///loop over all the particles and their energies
    ///to calculate invariant jet mass

    for(unsigned int k=0; k < p.size() ; k++){

      en += energy[k];
      psum += p[k];
    }
    pn = mod(psum);

    ///set value of m equal to the invariant jet mass of the 
    ///event stored in array _calcInvariantJetMass
    m = ((en * en) - (pn * pn)) / (4*(en * en));
  }       

  
  
  
  
  // Do the full calculation
  void InvariantJetMass::_calcInvariantJetMass(const vector<FourMomentum>& fsmomenta) {
       
    // Clear the caches
    _invariantjetmasss.clear();
    
    
    // If there are fewer than 2 visible particles, we can't do much
    if (fsmomenta.size() < 2) {
      for (int i = 0; i < 3; ++i) {
	_invariantjetmasss.push_back(-1);
      }
      return;
    }          
    
    // Temporary variables for calbs

    double val = 0.;
    
    // Get invariantjetmass
    _calcM(fsmomenta, val);

    //store the invariant jet mass to be called by cparam():
    _invariantjetmasss.push_back(val);

    
  }
  
  
  
}