#ifndef __JAANET__JLORENTZBOOST__
#define __JAANET__JLORENTZBOOST__

#include <vector>

#include "km3net-dataformat/offline/Vec.hh"
#include "km3net-dataformat/offline/Evt.hh"
#include "km3net-dataformat/offline/Trk.hh"
#include "km3net-dataformat/offline/Hit.hh"

#include "JLang/JException.hh"

#include "JGeometry3D/JVector3D.hh"
#include "JGeometry3D/JVertex3D.hh"
#include "JGeometry3D/JTrack3E.hh"

#include "JAAnet/JAAnetToolkit.hh"

#include "JPhysics/JConstants.hh"

#include "Math/Vector4D.h"
#include "Math/GenVector/Boost.h"
#include "Math/GenVector/LorentzVector.h"


/**
 * \file
 *
 * Auxiliary methods for calculating Lorentz boosts.
 * \author bjjung
 */

namespace JAANET {}
namespace JPP { using namespace JAANET; }

namespace JAANET {

  using ROOT::Math::Boost;
  using ROOT::Math::LorentzVector;


  /**
   * Auxiliary class for performing Lorentz boosts.
   */
  struct JLorentzBoost :
    public Boost
  {
    /** 
     * Default constructor.
     */
    JLorentzBoost()
    {}

    
    /**
     * Copy constructor.
     *
     * \param  boost         Lorentz boost
     */
    JLorentzBoost(const Boost& boost) :
      Boost(boost)
    {}


    /**
     * Retrieve gamma factor.
     *
     * \return               gamma factor
     */
    double getGamma() const
    {
      const double beta2 = this->BetaVector().Mag2();
      
      return 1.0 / sqrt(1.0 - beta2);
    }
    

    /**
     * Lorentz boost operator.
     *
     * \param  x4            4-vector
     * \return               boosted 4-vector
     */
    template<class CoordSystem>
    LorentzVector<CoordSystem> operator()(const LorentzVector<CoordSystem>& x4) const
    {
      return static_cast<const Boost&>(*this)(x4);
    }
    

    /**
     * Lorentz boost operator.
     *
     * \param  vertex        event vertex
     */
    void operator()(JVertex3D& vertex) const
    {
      using namespace JPP;
      using namespace ROOT::Math;
      
      const XYZTVector x0(vertex.getX(), vertex.getY(), vertex.getZ(), vertex.getT() * getSpeedOfLight());
      const XYZTVector x1 = (*this)(x0);

      const JPosition3D pos1(x1.X(), x1.Y(), x1.Z());

      vertex.setPosition(pos1);
      vertex.setT(x1.T() / getSpeedOfLight());
    }

    
    /**
     * Lorentz boost operator.
     *
     * \param  track         track
     * \param  mass          track particle mass
     */
    void operator()(JTrack3E&    track,
		    const double mass) const
    {
      using namespace JPP;
      using namespace ROOT::Math;
      
      const double Ekin = getKineticEnergy(track.getE(), mass);
      
      const XYZTVector x0(track.getX(), track.getY(), track.getZ(), track.getT() * getSpeedOfLight());
      const XYZTVector x1 = (*this)(x0);

      const PxPyPzEVector p0(Ekin * track.getDX(), Ekin * track.getDY(), Ekin * track.getDZ(), track.getE());
      const PxPyPzEVector p1 = (*this)(p0);

      const JVector3D pos1(x1.X(), x1.Y(), x1.Z());
      const JVersor3D dir1(p1.X(), p1.Y(), p1.Z());

      track.setPosition (pos1);
      track.setDirection(dir1);
      track.setT(x1.T());
      track.setE(p1.E());
    }
    

    /**
     * Lorentz boost operator.
     * N.B.: Only initial and final state tracks are boosted.
     *
     * \param  track         track
     */
    void operator()(Trk& track) const
    {
      using namespace std;
      using namespace JPP;
      using namespace ROOT::Math;

      // Boost momentum 4-vector

      if (is_finalstate(track) || is_initialstate(track)) {
      
	const Vec Ekin = (track.status != TRK_ST_ININUCLEI ?
			  getKineticEnergy(track) * track.dir : Vec(0.0, 0.0, 0.0));

	const PxPyPzEVector p0(Ekin.x, Ekin.y, Ekin.z, track.E);
	const PxPyPzEVector p1 = (*this)(p0);
      
	track.E   = p1.E();
	track.dir = Vec(p1.Px(), p1.Py(), p1.Pz());
	if (track.dir.len() > 0) { track.dir.normalize(); }

	// Boost track vertex and compute new track length,
	// assuming a relativistic particle traveling at the speed of light

	JVertex3D x0(getPosition(track.pos), track.t);         // Vertex 
	JVertex3D x1(getPosition(track.pos + track.len * track.dir),
		     track.t + track.len / getSpeedOfLight()); // Vertex + travel distance
      
	(*this)(x0); // Boost vertices
	(*this)(x1);

	track.t   = x0.getT() / getSpeedOfLight();
	track.pos = Vec(x0.getX(), x0.getY(), x0.getZ());

	track.len = x0.getDistance(x1);
      }
    }


    /**
     * Lorentz boost operator.
     * N.B.: The time-over-threshold is not boosted.
     *
     * \param  hit           hit
     */
    void operator()(Hit& hit) const
    {
      using namespace std;
      using namespace JPP;

      const double dt = hit.t - hit.tdc;
      
      JVertex3D vertex(getPosition(hit.pos), hit.t);
      (*this)(vertex);

      hit.pos = Vec(vertex.getX(), vertex.getY(), vertex.getZ());
      hit.t   = vertex.getT();

      hit.tdc = hit.t - getGamma() * dt;
    }


    /**
     * Lorentz boost operator.
     *
     * \param  __begin       begin of data
     * \param  __end         end of data
     */
    template<class T>
    void operator()(T __begin, T __end) const
    {
      using namespace std;
      
      for (T i = __begin; i != __end; ++i) {
	(*this)(*i);
      }
    }


    /**
     * Lorentz boost operator.
     * N.B.: Intermediate state tracks as well as\n
     *       the event timestamps and MC-time variable are not boosted.
     *
     * \param  event         event
     */
    void operator()(Evt& event) const
    {
      (*this)(event.mc_trks.begin(), event.mc_trks.end());
      (*this)(event.mc_hits.begin(), event.mc_hits.end());
      (*this)(event.trks.begin(),    event.trks.end());
      (*this)(event.hits.begin(),    event.hits.end());      
    }
  };


  /**
   * Get Lorentz boost to the Center of Mass (COM) frame for a given neutrino interaction.
   *
   * \param  event           event
   */
  JLorentzBoost getBoostToCOM(const Evt& event)
  {
    using namespace JPP;
    
    if (has_neutrino(event)) {

      const Trk& nu = get_neutrino(event);

      const Vec beta = -nu.E / getE0(event) * nu.dir;

      return Boost(beta.x, beta.y, beta.z);
      
    } else {
      THROW(JValueOutOfRange, "getBoostToCOM(): Given event does not correspond to a neutrino interaction.");
    }
  }


  /**
   * Boost event to the Center of Mass (COM) frame.
   *
   * \param  event           event
   */
  void boostToCOM(Evt& event)
  {
    const JLorentzBoost boost = getBoostToCOM(event);
    boost(event);
  }
}

#endif