#ifndef H_J_SCATTERING_MODEL
#define H_J_SCATTERING_MODEL

/**
 * \author mjongen
 */

// c++ standard library includes
#include<iostream>
#include<iomanip>
#include<cmath>
#include<cstdlib>

// root includes
#include "TObject.h"
#include "TH1F.h"
#include "TH2F.h"
#include "TRandom3.h"

// JPP includes
// these include files are hidden from rootcint
#include "JROOT/JRoot.hh"
#include "JMarkov/JPhotonPath.hh"
#include "JGeometry3D/JAxis3D.hh"
#include "JGeometry3D/JPosition3D.hh"
#include "JGeometry3D/JDirection3D.hh"

#ifndef __CINT__
  #include "JPhysics/KM3NeT.hh"
#endif

namespace JMARKOV {

  // for rootcint, all we need is this class declaration
  #ifdef __CINT__
    class JPhotonPath ;
  #endif
}

namespace JMARKOV {
  using namespace JGEOMETRY3D ;

  /**
     Virtual base class for a scattering model.
   **/
  class JScatteringModel {
    public :

    // constructor
    JScatteringModel() {}

    // destructor
    virtual ~JScatteringModel() {}

    /**
       Return the probability density as a function of cos(theta)
       
       dP / dOmega = dP / dcosTheta dPhi 

       to scatter in a given direction. Theta=0 is forward scattering.
    **/
    virtual double getScatteringProbability( double ct ) = 0 ;

    /**
       Return a randomly generated direction according to the scattering
       probability distribution.

       This uses gRandom.
     **/
    virtual JVersor3D generateDirection() = 0 ;

    // return the scattering length in [m]
    virtual double getScatteringLength() { return lambda_scat ; }

    protected :

    double lambda_scat ; // scattering length in [m]
  } ;

  /**
     Virtual base class for a light source.
   **/
  class JSourceModel {
    public: 

    JSourceModel() {}

    virtual ~JSourceModel() {}

    /**
       Return the probability density 
       
       dP / dOmega = dP / dCosTheta dPhi 

       that a photon from this source is emitted in a given direction, 
       given that a photon is emitted.
    **/
    virtual double getEmissionProbability( JVersor3D dir ) = 0 ;

    /**
       Return a randomly generated direction according to the emission
       distribution.
       
       This uses gRandom.
     **/
    virtual JVersor3D generateDirection() = 0 ;

    void setPosition( JPosition3D& _pos ) { pos = _pos ; }

    const JPosition3D& getPosition() const { return pos ; }

    protected:

    JPosition3D pos ;
  } ;

  /**
     Implementation of the JSourceModel class that represents an
     isotropic source.
   **/
  class JIsotropicSource : public JSourceModel {
    public : 

    JIsotropicSource() {}
    
    double getEmissionProbability( JVersor3D dir ) {
      return 1.0/(4.0*M_PI) ;
    }
    
    JVersor3D generateDirection() {
      double dx, dy, dz ;
      gRandom->Sphere(dx,dy,dz,1) ;
      return JVersor3D(dx,dy,dz) ;
    }
  } ;

  /**
     Implementation of the JSourceModel class that represents a
     directed source with a flat intensity distribution.

     When _cutSurface = true, photons that would be emitted into 
     the surface (defined by the surface normal) will not be emitted.
   **/
  class JDirectedSource : public JSourceModel {
    public : 

    /**
       Constructor. 
       _source_dir is the source direction
       _alpha      is the opening angle in radians
       
       note that _alpha should be > -1 
     **/
    JDirectedSource( JVersor3D _source_dir, double _alpha, bool _cutSurface=false, const JVersor3D _surface_normal=JVersor3D() ) :
      source_dir(_source_dir), ctmin(cos(0.5*_alpha)), alpha(_alpha), cutSurface(_cutSurface), surface_normal(_surface_normal) {
      // we have to find out which fraction of photons would be emitted into the surface, given
      // the direction and opening angle
      // as far as I can tell, this is a non-trivial integral, so I solve it numerically

      // delta is the angle of the beam direction over the surface 
      // (negative means the beam points into the surface)
      delta = 0.5*M_PI - acos(surface_normal.getDot(source_dir)) ;
      omega = 2.0*M_PI*(1.0-ctmin) ; // space angle of beam ("spherical cap")

      if( cutSurface ) {
	if( delta>0.0 && delta>0.5*_alpha )  {
	  // no part of the beam points into the surface
	  //fraction_into_surface = 0.0 ;
	} else if( delta<=0 && fabs(delta)>=0.5*alpha ) {
	  // the beam points completely into the surface, throw an error
	  //fraction_into_surface = 1.0 ;
	  omega = 0.0 ;
	  cerr << "FATAL ERROR in JDirectedSource: beam points completely into the surface." << endl ;
	  exit(1) ;
	} else {
	  // part of the beam points into the surface, which effectively reduces
	  // the space angle omega
	  // here we numerically compute omega

	  // only at theta=delta to alpha/2 is part of the beam absorbed
	  // we calculate the space angle of the absorbed beam part
	  double omega_absorbed = 0.0 ;
	  const int __nct = 100000 ;
	  const double __ctmin = cos(0.5*alpha) ;
	  const double __ctmax = cos(delta) ;
	  const double __dct   = (__ctmax - __ctmin)/__nct ;
	  for( int i=0; i<__nct; ++i ) {
	    double __ct = __ctmin + (i+0.5) * __dct ;
	    double __theta = acos(__ct) ;
	    double __phirange = 2.0*(M_PI-getPhiMax(__theta)) ; 
	    omega_absorbed += __phirange ;
	  }
	  omega_absorbed *= __dct ;
	  omega -= omega_absorbed ;

	  // if the beam direction is into the surface, this part gets absorbed completely
	  if( delta<0 ) {
	    omega -= 2.0 * M_PI * (1.0-cos(delta)) ;
	  }
	}
      }
    }
    
    /// return the opening angle in radians
    double getOpeningAngle() const { return alpha ; }

    /// return the source direction
    JVersor3D getDirection() const { return source_dir ; }

    double getEmissionProbability( JVersor3D dir ) {
      double ct = dir.getDot(source_dir) ;
      if( ct>ctmin ) {
	if( cutSurface && dir.getDot(surface_normal)<0 ) return 0.0 ;
	return 1.0/omega ;
      }
      return 0.0 ;
    }

    JVersor3D generateDirection() {
      while(true) {
	double phi = gRandom->Uniform(0,2*M_PI) ;
	double ct  = gRandom->Uniform(ctmin,1) ;
	JDirection3D _dir( JAngle3D(acos(ct),phi) ) ;      
	JRotation3D R( source_dir ) ; 
	JVersor3D dir(_dir.rotate_back(R)) ;
	if( !cutSurface ) return dir ;
	if( dir.getDot(surface_normal)>0 ) return dir ;
      }
    }

    protected :

    /**
       For a given angle theta with respect to the beam direction, and a given surface
       normal, we define phi as the rotation angle around the beam, where phi=0
       corresponds to the direction sticking out over the surface a much as possible.

       This function returns 'phi max', which is the largest value of phi for which the
       direction does not point into the surface.

    **/
    double getPhiMax( double theta ) {
      if( delta>0 ) {
	if( theta<delta ) return M_PI ;
	if( theta+delta > M_PI ) return 0.0 ;
      }
      if( delta<0 ) {
	if( theta<fabs(delta) ) return 0.0 ;
	if( theta+fabs(delta) > M_PI ) return M_PI ;
      }
      double phimax = acos(-tan(delta)/tan(theta)) ;
      return phimax ;
    }
    

    JVersor3D source_dir ;
    double ctmin ;
    double alpha ;
    bool cutSurface ;
    JVersor3D surface_normal ;
    double delta ;
    //double fraction_into_surface ; // the fraction of photons that would be emitted into the surface
    double omega ; // integrated phase angle of photon emission directions
  } ;


  /**
     Virtual base class for a light detector ("photon target").
  **/
  class JTargetModel {
    public: 

    JTargetModel() {
      r = 0 ; // infinitely small target by default
    }

    virtual ~JTargetModel() {}

    /**
       Return the efficiency, which is defined as the probability that
       a photon with the given direction hitting the target will be 
       registered. 

       Note that we assume by convention that the direction is the 
       PHOTON direction, NOT the direction that you would see the 
       photon coming from!
       
       By convention the highest efficiency is 1.
    **/
    virtual double getEfficiency( JVersor3D dir ) const = 0 ;

    /**
       Return the effective area, i.e. the integral over the whole 
       surface of the target, weighted by the efficiency.
     **/
    virtual double getEffectiveArea() { return 0.0 ; }

    // set the target radius in [m]
    void setRadius( double _r ) { r = _r ; }

    // get the target radius in [m]
    double getRadius() const { return r ; }

    void setPosition( JPosition3D& _pos ) { axis = JAxis3D(_pos,axis.getDirection()) ; }
    const JPosition3D& getPosition() const { return axis.getPosition() ; }

    void setDirection( JVersor3D& _dir ) { axis = JAxis3D(axis.getPosition(),_dir) ; }
    const JVersor3D& getDirection() const { return axis.getDirection() ; }

    protected:

    double r ;
    JAxis3D axis ;
  } ;

  /**
     Implementation of the JTargetModel class that represents a
     single PMT on a DOM
   **/
  class JPMTTarget : public JTargetModel {
    public :
    
    /**
       Constructor

       _r is DOM radius, default = 0.2159 [m] (=17 inch diameter glass sphere)

       _alpha is the opening angle in radians
    **/
    JPMTTarget( JAxis3D _axis, double _alpha, double _r=0.2159 ) {
      setRadius(_r) ;
      axis = _axis ;
      alpha = _alpha ;

      // minimal dot product of PMT direction vs hit direction on DOM
      ctmin = cos(0.5*_alpha) ;
    }

    /// return the PMT opening angle in radians
    double getOpeningAngle() const { return alpha ; }

    double getEfficiency( JVersor3D dir ) const {
      const double ct = axis.getDirection().getDot(dir) ;
      if( ct > ctmin ) {
	return 1.0 ;
      } else {
	return 0.0 ;
      }
    }

    double getEffectiveArea() const { 
      return 2.0*r*r*M_PI*(1-ctmin) ; 
    }
    
    double ctmin ;

    protected :

    double alpha ;
  } ;
  
  /**
     Implementation of the JTargetModel class that represents a
     spherically symmetric target
   **/
  class JIsotropicTarget : public JTargetModel {
    public : 

    JIsotropicTarget() {}
    
    double getEfficiency( JVersor3D dir ) const { return 1.0 ; }
  } ;

  /**
     Implementation of the JTargetModel class that represents a
     directed target with a cos(theta) acceptance.
   **/
  class JCosineTarget : public JTargetModel {
    public : 

    /**
       Constructor. 
       _target_dir        is the direction where the photon has
                          the highest chance to be detected.
     **/
    JCosineTarget( JVersor3D _target_dir ) :
      target_dir(_target_dir) {}
    
    double getEfficiency( JVersor3D dir ) const {
      return max( dir.getDot(target_dir), 0.0 ) ;
    }

    protected :

    JVersor3D target_dir ;
  } ;


  /**
     Implementation of the JScatteringModel interface with scattering
     according to the Henyey-Greenstein function.
   **/
  class JHenyeyGreensteinScattering : public JScatteringModel {

   public:

    /**
       Constructor.

       g is a parameter that determines how forward the scattering is. 
       g = 0 corresponds to isotropic scattering
       g = 1 corresponds to fully forward scattering
     **/
    JHenyeyGreensteinScattering( double _lambda_scat, double _g ) {
      lambda_scat  = _lambda_scat ;
      g = _g ;
    }

    double getScatteringProbability( double ct ) {
      //double ct = cos(dir.getTheta()) ;
      double den = 1 + g*g - 2*g*ct ;
      return (1-g*g)/(4*M_PI*den*sqrt(den)) ;
    }

    JVersor3D generateDirection() {
      double phi = gRandom->Uniform(0,2*M_PI) ;
      double ct ;
      if( g>0 ) { 
	double xi  = gRandom->Uniform(0,1) ;
	ct  = 1.0/(2*g) * ( 1 + g*g - (1-g*g)*(1-g*g)/( (1-g+2*g*xi)*(1-g+2*g*xi) ) ) ;
      } else {
	ct = gRandom->Uniform(-1,1) ;
      }
      return JDirection3D( JAngle3D( acos(ct),phi) ) ;
    }

   protected :

    double g ;
 
  } ;

  /**
     Implementation of the JScatteringModel interface with
     isotropic scattering
   **/
  class JIsotropicScattering : public JScatteringModel {

   public:

    /**
       Constructor.
     **/
    JIsotropicScattering( double _lambda_scat ) {
      lambda_scat  = _lambda_scat ;
    }

    double getScatteringProbability( double ct ) {
      return 1.0/(4.0*M_PI) ;
    }

    JVersor3D generateDirection() {
      double phi = gRandom->Uniform(0,2*M_PI) ;
      double ct  = gRandom->Uniform(-1,1) ;
      return JDirection3D( JAngle3D( acos(ct),phi) ) ;
    }
  } ;

  /**
     Implementation of the JScatteringModel interface with
     Rayleigh scattering
   **/
  class JRayleighScattering : public JScatteringModel {

   public:

    /**
       Constructor.

       note that one should never set a<0
     **/
    JRayleighScattering( double _lambda_scat, double _a ) {
      lambda_scat  = _lambda_scat ;
      a = _a ;
      if( fabs(a)<0 ) {
	cerr << "Fatal error in initialization of JRayleighScattering: a<0 !" << endl ;
	exit(1) ;
      }
    }

    double getScatteringProbability( double ct ) {
      //double ct = cos(dir.getTheta()) ;
      return  (1+a*ct*ct)/(4.0*M_PI*(1+a/3.0)) ;
    }

    JVersor3D generateDirection() {
      double phi = gRandom->Uniform(0,2*M_PI) ;
      double ct ;
      if( a>0 ) {
	double xi  = gRandom->Uniform(0,1) ;
	double d0 = -9/a ;
	double d1 = 27.0*(1-2*xi)*(3+a)/a ;
	double C  = pow( 0.5*(d1 + sqrt(d1*d1-4*d0*d0*d0) ), 1.0/3.0 ) ;
	ct = -1.0/(3.0)*(C+d0/C) ;
      } else {
	ct = gRandom->Uniform(-1,1) ;
      }
      return JDirection3D( JAngle3D( acos(ct),phi) ) ;
    }

  protected:

    double a ;
  } ;

  /**
     Implementation of the JScatteringModel interface with
     that combines two scattering models into one effective model.
   **/
  class JCombinedScattering : public JScatteringModel {

   public:

    /**
       Constructor.

       Note that this only copies only the pointers to the JScatteringModel 
       instances. So if _sm1 or _sm2 are deleted or changed, it will also
       affect the behaviour of this instance.
     **/
    JCombinedScattering( JScatteringModel* _sm1, JScatteringModel* _sm2 ) : sm1(_sm1), sm2(_sm2) { }

    virtual double getScatteringLength() {
      // calculate effective scattering length
      double l1 = sm1->getScatteringLength() ;
      double l2 = sm2->getScatteringLength() ;
      return 1.0 / ( 1.0/l1 + 1.0/l2 ) ;
    }

    double getScatteringProbability( double ct ) {
      double val1 = sm1->getScatteringProbability(ct) ;
      double val2 = sm2->getScatteringProbability(ct) ;
      double l1 = sm1->getScatteringLength() ;
      double l2 = sm2->getScatteringLength() ;
      double lambda = getScatteringLength() ; 
      // effective scattering probability
      return lambda * ( val1/l1 + val2/l2 ) ;
    }

    JVersor3D generateDirection() {
      double l1 = sm1->getScatteringLength() ;
      double l2 = sm2->getScatteringLength() ;
      // probability to scatter with process 1
      double P = l2/(l1+l2) ; 
      if( gRandom->Uniform(0,1)<P ) {
	return sm1->generateDirection() ;
      } else { 
	return sm2->generateDirection() ;
      }
    }

  protected:

    double a ;
    JScatteringModel* sm1 ;
    JScatteringModel* sm2 ;
  } ;

  /**
     Return the probability density for a photon path with the given
     ingredients. These are
     - a model of the source direction distribution
     - a scattering model
     - a target model
     - the absorption length (note that this is allowed to be +INFTY)

     Multiply by dV1, dV2, etc. and the target cross-section sigma 
     to get a probability.

     The first (last) vertex of the photon path is used as the source 
     (target) position.
  **/
  double getPhotonPathProbabilityDensity( JPhotonPath& p, JSourceModel* src, JScatteringModel* sm, JTargetModel* trg, double lambda_abs ) {
    //const bool verbose = true ;

    // calculate inverse scattering and absorption lengths
    const double lambda_scat_inv = 1.0/sm->getScatteringLength() ;
    double lambda_abs_inv = 0 ;
    // if the absorption length is not infinite
    if( lambda_abs == lambda_abs ) lambda_abs_inv = 1.0/lambda_abs ;

    double rho = 1 ;
    const int nvert = p.n   ; // number of vertices in the path
    const int nscat = p.n-2 ; // number of scattering vertices in the path

    //if( verbose ) cout << "$ rho = " << rho << " (before calculating path segment lengths)" << endl ;

    // calculate path segment lengths
    // index i corresponds to the path segment connecting 
    // vertex i and vertex i+1
    double Ltotal = 0 ;
    vector<double> lengths(nvert-1) ; 
    for( int i=0 ; i<nvert-1 ; ++i ) {
      lengths[i] = (p[i+1]-p[i]).getLength() ;
      Ltotal += lengths[i] ;
    }

    //if( verbose ) cout << "$ Total length = " << Ltotal << " m" << endl ;

    // cosine of scattering angles (angle between previous segment and new segment)
    // index i is the cosine of the angle of segment i+1 w.r.t segment i
    vector<double> cts(nscat) ;

    // scattering angles
    for( int i=0 ; i<nscat ; ++i ) {
      JDirection3D seg1(p[i+1]-p[i]) ;
      JDirection3D seg2(p[i+2]-p[i+1]) ;
      cts[i] = seg1.getDot(seg2) ;
    }

    // term for absorption length
    // (=probability to not be absorbed within Ltotal)
    rho *= exp(-Ltotal*lambda_abs_inv ) ;

    //if( verbose ) cout << "$ rho = " << rho << " (after absorption length)" << endl ;

    // term for emission profile
    rho *= src->getEmissionProbability( JDirection3D(p[1]-p[0]) ) ;

    //if( verbose ) cout << "$ rho = " << rho << " (after emission probability)" << endl ;

    // term for target efficiency
    if( trg->getRadius()>0 ) {
      // finite target (spherical), the efficiency is a measure of how good the
      // target is at a given spot
      rho *= trg->getEfficiency( JDirection3D(p[nscat+1]-trg->getPosition()) ) ;
      // this factor accounts for the angle of incidence
      double ct = JDirection3D(p[nscat]-p[nscat+1]).getDot( JDirection3D(p[nscat+1]-trg->getPosition()) ) ;
      rho *= max(0.0,ct) ;

      //if( verbose ) cout << "$ rho = " << rho << " (after target efficiency and angle of incidence)" << endl ;

      // check whether the photon path hits the sphere
      // we do not need to consider the last vertex, because if that intersects the sphere
      // the angle of incidence factor above would be 0.
      if( p.n>2 ) {
	JPhotonPath pshort(p) ;
	pshort.resize( pshort.size()-1 ) ;
	--pshort.n ;
	if( pshort.hitsSphere(trg->getPosition(),trg->getRadius()) ) {
	  /*cout << "Path hits sphere at " 
	       << JPosition3D(pshort.getSphereHitPosition(trg->getPosition(),trg->getRadius())-trg->getPosition()) 
	       << endl ;*/
	  rho *= 0 ;
	}
      }
    } else {
      rho *= trg->getEfficiency( JDirection3D(p[nscat+1]-p[nscat]) ) ;
    }

    //if( verbose ) cout << "$ rho = " << rho << " (after target efficiency)" << endl ;

    // terms for scattering directions
    for( int i=0 ; i<nscat ; ++i ) {
      rho *= sm->getScatteringProbability(cts[i]) ;
    }

    //if( verbose ) cout << "$ rho = " << rho << " (after scattering directions)" << endl ;

    // phase space terms for scattering to each of the scattering 
    // vertices
    // this is the probability/dV to scatter into a volume element 
    // dV at a given distance
    // (although it is missing the direction-dependent term, which is
    // the emission probability or scattering probability)
    for( int i=0; i<nscat; ++i ) {
      // reminder: lengths[i] is the length from vertex i to vertex i+1
      rho *= lambda_scat_inv / (lengths[i]*lengths[i]) * exp(-lengths[i]*lambda_scat_inv) ;
    }

    //if( verbose ) cout << "$ rho = " << rho << " (after scattering phase space term)" << endl ;

    // phase space term for scattering to the target
    // this is the probability/dsigma to hit a target at a given
    // distance, without scattering along the way
    rho *= exp(-lengths[nscat]*lambda_scat_inv) / ( lengths[nscat]*lengths[nscat] ) ;

    //if( verbose ) cout << "$ rho = " << rho << " (after target phase space term)" << endl ;

    // check for NaN
    if( rho!=rho ) {
      cerr << "NaN for probability density" << endl ;
      cerr << "rho = " << rho << endl ;
      exit(1) ;
    }

    //if( verbose ) cout << "$ rho = " << rho << endl ;

    return rho ;
  }

}

#endif