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


namespace Rivet {
  
  
  void Broadening::calc(const FinalState& fs) {
    calc(fs.particles());
  }
  
  void Broadening::calc(const vector<Particle>& fsparticles) {
    vector<Vector3> threeMomenta;
    threeMomenta.reserve(fsparticles.size());
    foreach (const Particle& p, fsparticles) {
      const Vector3 p3 = p.momentum().vector3();
      threeMomenta.push_back(p3);
     }
    _calcBroadening(threeMomenta);
  }
  
  void Broadening::calc(const vector<FourMomentum>& fsmomenta) {
    vector<Vector3> threeMomenta;
    threeMomenta.reserve(fsmomenta.size());
    foreach (const FourMomentum& v, fsmomenta) {
      threeMomenta.push_back(v.vector3());
    }
    _calcBroadening(threeMomenta);
  }
  
  void Broadening::calc(const vector<Vector3>& fsmomenta) {
    _calcBroadening(fsmomenta);
  }
  
  
  /////////////////////////////////////////////////
  
  
  
  inline bool mod2Cmp(const Vector3& a, const Vector3& b) {
    return a.mod2() > b.mod2();
  }
  
  
  // Do the general case broadening calculation
  void _calcB(const vector<Vector3>& momenta, double& t, Vector3& taxis) {
    /* This function implements the iterative algorithm as described in the
     * Pythia manual. We take eight (four) different starting vectors
      * constructed from the four (three) leading particles to make sure that
      * we don't find a local maximum.
      */
    vector<Vector3> p = momenta;
    assert(p.size() >= 3);
    unsigned int n = 3;
    if (p.size() == 3) n = 3;
    vector<Vector3> tvec;
    vector<Vector3> pvec;
    Vector3 pveci;
    vector<double> tval;
    std::sort(p.begin(), p.end(), mod2Cmp);
    Vector3 foo(0,0,0);
    int sign;
    for (unsigned int i=0 ; i<pow(2,n-1) ; i++) {
      // Create an initial vector from the leading four jets
      sign=i;
      for (unsigned int k=0 ; k<n ; k++) {
	(sign%2)==1 ? foo+=p[k] : foo-=p[k];
	sign/=2;
      }
      foo=foo.unit();
      
      // Iterate
      double diff=999.;
      while (diff>1e-5) {
	Vector3 foobar(0,0,0);
	for (unsigned int k=0 ; k<p.size() ; k++) 
	  foo.dot(p[k])>0 ? foobar+=p[k] : foobar-=p[k];
	diff=(foo-foobar.unit()).mod();
	foo=foobar.unit();
      }
      
      // Calculate the thrust value for the vector we found
      t=0.;
      for (unsigned int k=0 ; k<p.size() ; k++) 
	t+=fabs(foo.dot(p[k]));
      // Store everything
      tval.push_back(t);
      tvec.push_back(foo);
      
    }
    
    // Pick the solution with the largest thrust
    t=0.;
    for (unsigned int i=0 ; i<tvec.size() ; i++) 
      if (tval[i]>t){
	t=tval[i];
	taxis=tvec[i];
      }
    
    
    // Calculate the broadening value for the vector we found
    t=0.;
    for (unsigned int k=0 ; k<p.size() ; k++) {
      t+=mod(taxis.cross(p[k]));//Only differance from Thrust.cc is a cross product instead of a dot product
    }
    
  }
  
  
  
  
  // Do the full calculation
  void Broadening::_calcBroadening(const vector<Vector3>& fsmomenta) {
    // Make a vector of the three-momenta in the final state
    double momentumSum(0.0);
    foreach (const Vector3& p3, fsmomenta) {
      momentumSum += mod(p3);
    }
    getLog() << Log::DEBUG << "Number of particles = " << fsmomenta.size() << endl;
    
    
    // Clear the caches
    _broadenings.clear();
    _broadeningAxes.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) {
	_broadenings.push_back(-1);
	_broadeningAxes.push_back(Vector3(0,0,0));
      }
      return;
    }
    
    
    // Handle special case of broadening = 1 if there are only 2 particles
    if (fsmomenta.size() == 2) {
      Vector3 axis(0,0,0);
      _broadenings.push_back(0.0);
      _broadenings.push_back(0.0);
      _broadenings.push_back(0.0);
      axis = fsmomenta[0].unit();
      if (axis.z() < 0) axis = -axis;
      _broadeningAxes.push_back(axis);
      /// @todo Improve this --- special directions bad...
      /// (a,b,c) _|_ 1/(a^2+b^2) (b,-a,0) etc., but which combination minimises error?
      if (axis.z() < 0.75)
	_broadeningAxes.push_back( (axis.cross(Vector3(0,0,1))).unit() );
       else
         _broadeningAxes.push_back( (axis.cross(Vector3(0,1,0))).unit() );
       _broadeningAxes.push_back( _broadeningAxes[0].cross(_broadeningAxes[1]) );
      return;
    }
    
    
    
    // Temporary variables for calbs
    Vector3 axis(0,0,0);
    double val = 0.;
    
    // Get broadening
    _calcB(fsmomenta, val, axis);
    getLog() << Log::DEBUG << "Mom sum = " << momentumSum << endl;
    _broadenings.push_back(val / momentumSum);
    // Make sure that broadening always points along the +ve z-axis.
    if (axis.z() < 0) axis = -axis;
    axis = axis.unit();
    getLog() << Log::DEBUG << "Axis = " << axis << endl;
    _broadeningAxes.push_back(axis);
    
  }
  
  
  
}