/// -*- C++ -*-
#include "Rivet/Rivet.hh"
#include "ThrustGamma.hh"
#include "Rivet/Tools/Logging.hh"
#include "Rivet/Cmp.hh"
#include "Rivet/Projections/DISKinematics.hh"
#include "Rivet/Math/Constants.hh"
#include "Rivet/Tools/ParticleIdUtils.hh"
#include "math.h"
#include "Rivet/Projection.hh"
#include "Rivet/Projections/AxesDefinition.hh"
#include "Rivet/Projections/FinalState.hh"
#include "Rivet/Event.hh"
#include "Rivet/Particle.hh"
#include "Rivet/Projections/DISKinematics.hh"
#include "Rivet/Projections/DISLepton.hh"
#include "Rivet/Projections/Beam.hh"

namespace Rivet {
    
  int ThrustGamma::compare(const Projection& p) const {
    return mkNamedPCmp(p, "Kinematics");
  }

  ThrustGamma::ThrustGamma(const DISKinematics& kinematicsp, const FinalState& fs){
    setName("ThrustGamma");
    addProjection(kinematicsp, "Kinematics");
    addProjection(Beam(), "Beam");
    addProjection(DISLepton(), "Lepton");
    addProjection(fs, "FS");

  }

  void ThrustGamma::project(const Event& e) {
    const DISKinematics& diskin = applyProjection<DISKinematics>(e, "Kinematics");
    const LorentzTransform breitboost = diskin.boostBreit();
    const DISLepton& dislep = diskin.applyProjection<DISLepton>(e, "Lepton");
    const GenParticle& dislepGP = dislep.out().genParticle();
    const FinalState& fs = dislep.applyProjection<FinalState>(e, "FS");
    Particle _inHadron;
    LorentzTransform _hcm;
    
    //This is mostly to find the breit shifted PGamma to make sure that the 
    //final state particles I store are in the chosen hemisphere:

    // Identify beam hadron
    const ParticlePair& inc = applyProjection<Beam>(e, "Beam").beams();
    bool firstIsHadron  = PID::isHadron(inc.first.pdgId());
    bool secondIsHadron = PID::isHadron(inc.second.pdgId());
    if (firstIsHadron && !secondIsHadron) {
      _inHadron = inc.first;
    } else if (!firstIsHadron && secondIsHadron) {
      _inHadron = inc.second;
    } else {
       //help!
      throw Error("DISKinematics projector could not find the correct beam hadron");
    }
    
    // several key DIS kinematics needed to work out the boost:
    FourMomentum pLepIn = dislep.in().momentum();
    FourMomentum pLepOut = dislep.out().momentum();
    FourMomentum pGamma = pLepIn - pLepOut;
    FourMomentum pHad = _inHadron.momentum();
    FourMomentum tothad = pGamma + pHad;
    double Q2 = -pGamma.mass2();
    double dx = diskin.x();

    // Calculate boost vector for boost into HCM-system:
    LorentzTransform tmp;
    tmp.setBoost(-tothad.boostVector());
    
    // Rotate so the photon is in x-z plane in HCM rest frame:
    FourMomentum pGammaHCM = tmp.transform(pGamma);
    tmp.preMult(Matrix3(Vector3::mkZ(), -pGammaHCM.azimuthalAngle()));
    pGammaHCM = tmp.transform(pGamma);
    
     // Rotate so the photon is along the positive z-axis:
    const double rot_angle = pGammaHCM.polarAngle() * (pGammaHCM.px() >= 0 ? -1 : 1);
    tmp.preMult(Matrix3(Vector3::mkY(), rot_angle));
    pGammaHCM = tmp.transform(pGamma);
 
    // Finally rotate so outgoing lepton at phi = 0:
    FourMomentum pLepOutHCM = tmp.transform(pLepOut);
    tmp.preMult(Matrix3(Vector3::mkZ(), -pLepOutHCM.azimuthalAngle()));
    _hcm = tmp;
    
    // Boost to Breit frame (use opposite convention for photon --- along *minus* z):
    tmp.preMult(Matrix3(Vector3::mkX(), PI));
    const double bz = 1 - 2*dx;
    LorentzTransform _breit = LorentzTransform(Vector3::mkZ() * bz).combine(tmp);
    
    //calculate the photons 3 momentum, aswell as the pre&post boost energys for cuts in the next stage:
    FourMomentum pGammaB = _breit.transform(pGamma);
    double bq= pGammaB.mod();
    double q= sqrt(Q2);
    Vector3 pgam = pGammaB.vector3();

    const vector<Particle> ps = applyProjection<FinalState>(e, "FS").particles();
    calcg(ps, pGammaB);//feeding the finalstate and fourmomentum to the next stage of the analysis.
  }
  

  
  void ThrustGamma::calcg(const vector<Particle>& fsparticles, FourMomentum pgamma) {
    vector<Vector3> threeMomenta;
    threeMomenta.reserve(fsparticles.size());
    foreach (const Particle& p, fsparticles) {
      const Vector3 p3 = p.momentum().vector3();
      threeMomenta.push_back(p3);
    }
    _calcgThrustGamma(threeMomenta,pgamma);
  }
  
  void ThrustGamma::calcg(const vector<FourMomentum>& fsmomenta, FourMomentum pgamma) {
    vector<Vector3> threeMomenta;
    threeMomenta.reserve(fsmomenta.size());
    foreach (const FourMomentum& v, fsmomenta) {
      threeMomenta.push_back(v.vector3());
    }
    _calcgThrustGamma(threeMomenta, pgamma);
  }
  
  void ThrustGamma::calcg(const vector<Vector3>& fsmomenta, FourMomentum pgamma) {
    _calcgThrustGamma(fsmomenta, pgamma);
  }
  
  
  /////////////////////////////////////////////////
  
  
  
  inline bool mod2Cmp(const Vector3& a, const Vector3& b) {
    return a.mod2() > b.mod2();
  }
  
  
  // Do the general case thrustgamma calculation
  void _calcBG(const vector<Vector3>& momenta, double& t, Vector3& taxis, FourMomentum pgamma) {
    /* 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<double> tval;
    std::sort(p.begin(), p.end(), mod2Cmp);
    for (unsigned int i=0 ; i<pow(2,n-1) ; i++) {
      // Create an initial vector from the leading four jets
      Vector3 foo(0,0,0);
      int sign=i;
      for (unsigned int k=0 ; k<n ; k++) {
	(sign%2)==1 ? foo+=p[k] : foo-=p[k];
	sign/=2;
      }
      foo=foo.unit();
            
      Vector3 pgam = pgamma.vector3();//make a 3vector from the fourmomentum of the photon

      // 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=pgam.unit(); // tried to avoid fiddling with the gubings of the analysis to essentialy just overode the various thrust axis so their all defined by the photon
      }
      
      // Calculate the thrustgamma value for the vector we found
      t=0.;
      for (unsigned int k=0 ; k<p.size() ; k++)
	t+=foo.dot(p[k]);
      
      // Store everything
      tval.push_back(t);
      tvec.push_back(foo);
    }
    
    // Pick the solution with the largest thrustgamma
    t=0.;
    for (unsigned int i=0 ; i<tvec.size() ; i++)
      if (tval[i]>t){
	t=tval[i];
	taxis=tvec[i];
      }
  }
  
  
  
  // Do the full calculation
  void ThrustGamma::_calcgThrustGamma(const vector<Vector3>& fsmomenta, FourMomentum pgamma) {
    // 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
    _thrustgammas.clear();
    _thrustgammaAxes.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) {
	_thrustgammas.push_back(-1);
	_thrustgammaAxes.push_back(Vector3(0,0,0));
      }
      return;
    }
    
    
    // Handle special case of thrustgamma = 1 if there are only 2 particles
    if (fsmomenta.size() == 2) {
      Vector3 axis(0,0,0);
      _thrustgammas.push_back(1.0);
      _thrustgammas.push_back(0.0);
      _thrustgammas.push_back(0.0);
      axis = fsmomenta[0].unit();
      if (axis.z() < 0) axis = -axis;
      _thrustgammaAxes.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)
	_thrustgammaAxes.push_back( (axis.cross(Vector3(0,0,1))).unit() );
       else
         _thrustgammaAxes.push_back( (axis.cross(Vector3(0,1,0))).unit() );
       _thrustgammaAxes.push_back( _thrustgammaAxes[0].cross(_thrustgammaAxes[1]) );
      return;
    }
    
    
    
    // Temporary variables for calcs
    Vector3 axis(0,0,0);
    double val = 0.;
    
    // Get thrustgamma
    _calcBG(fsmomenta, val, axis, pgamma);
    getLog() << Log::DEBUG << "Mom sum = " << momentumSum << endl;
    _thrustgammas.push_back(val / momentumSum);
    // Make sure that thrustgamma always points along the +ve z-axis.
    if (axis.z() < 0) axis = -axis;
    axis = axis.unit();
    getLog() << Log::DEBUG << "Axis = " << axis << endl;
    _thrustgammaAxes.push_back(axis);
    
    // Get thrustgamma major
    vector<Vector3> threeMomenta;
    foreach (const Vector3& v, fsmomenta) {
      // Get the part of each 3-momentum which is perpendicular to the thrustgamma axis
      const Vector3 vpar = dot(v, axis.unit()) * axis.unit();
      threeMomenta.push_back(v - vpar);
    }
    _calcBG(threeMomenta, val, axis, pgamma);
    _thrustgammas.push_back(val / momentumSum);
    if (axis.x() < 0) axis = -axis;
    axis = axis.unit();
    _thrustgammaAxes.push_back(axis);
    
    // Get thrustgamma minor
    if (_thrustgammaAxes[0].dot(_thrustgammaAxes[1]) < 1e-10) {
      axis = _thrustgammaAxes[0].cross(_thrustgammaAxes[1]);
      _thrustgammaAxes.push_back(axis);
      val = 0.0;
      foreach (const Vector3& v, fsmomenta) {
	val += fabs(dot(axis, v));
      }
      _thrustgammas.push_back(val / momentumSum);
    } else {
      _thrustgammas.push_back(-1.0);
      _thrustgammaAxes.push_back(Vector3(0,0,0));
    }
    
  }
  
  
  
}