// -*- 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 calculation
  void _calcB(const vector<FourMomentum>& momenta, double& m) {

    //create 2 vector objects, one storing the momenta 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.;
    double en=0.;
    double pn=0.; 
    //loop over all the particles and their energies to calculate the 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 in _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
    _calcB(fsmomenta, val);
    //store the invariant jet mass to be called by cparam():
    _invariantjetmasss.push_back(val);

    
  }
  
  
  
}