#ifndef H_MARKOV_PATH
#define H_MARKOV_PATH

/**
 * \author mjongen
 */

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

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

// JPP includes
#include "JMarkov/JScatteringModel.hh"

namespace JMARKOV {}
namespace JPP { using namespace JMARKOV; }

namespace JMARKOV {

  /**
     The JMarkovPathGenerator generates ensembles of photon paths using a 
     Markov Chain Monte Carlo (MCMC). 
     In the Markov chain, the number of scatterings does not change. 
     Therefore the ensembles have to be generated separately for each given 
     number of scatterings.

     The small modification of the path that is proposed in each step, is
     defined by the function 'randomPathChange'. This function is virtual
     so that it can be overrided in a derived class. 
     The initial and final vertex are fixed and one of the other vertices 
     is displaced in a random direction following an exponential distribution. 

     The source, target and scattering models are used to calculate the 
     probability density of the path. 

     Random numbers are drawn from gRandom.

     By default, the coordinates are remapped such that the 1/r^2 singularities
     in the path probability density disappear. This option can be turned off
     (but do not do it, unless you want to convince yourself that it is 
     necessary) with the setCoordinateRemapping method.
  **/
  class JMarkovPathGenerator {
    public:
    
    /// standard constructor
    JMarkovPathGenerator() {
      naccepted_r = -1 ;
      nrejected_r = -1 ;
      naccepted_t = -1 ;
      nrejected_t = -1 ;
      setTargetStepSize_deg(1.0) ;
      setRadialStepSize_m(1.0) ;
      setTangentialStepSize_deg(1.0) ;
      setCoordinateRemapping(true) ;
    }
  
    /** 
	Generate an ensemble of n paths with a fixed number of scatterings by 
	MCMC-sampling the given scattering model.
	The paths start at the source (a well-defined position) and end on the 
	target (either a fixed position or, in the case of a finite source, the
	surface of a sphere). 

	start_path is a photon path that is used a starting point for the MCMC.
	It has to be supplied by the user (due to the various source and target
	profiles it is not straightforward to generate one automatically) and 
	must have a path probability density != 0.
	The number of scatterings is also defined by the start_path.
    
	nsteps_burn_in is the number of steps taken for burn-in.

	nsteps_save determines the number of MCMC steps between the paths 
	that are saved to the ensemble.
    **/
    std::vector<JPhotonPath> generateEnsemble( int n, const JPhotonPath& start_path, JSourceModel* src, JScatteringModel* sm, JTargetModel* trg, double lambda_abs, int nsteps_burn_in, int nsteps_save ) ;

    /// get the number of Markov steps taken in the last call to generateEnsemble (after burn-in)
    int getNsteps() { return naccepted_r + naccepted_t + nrejected_r + nrejected_t ; }

    /// get the number of accepted steps taken during the last call to generateEnsemble (after burn-in)
    int getNacceptedSteps() { return naccepted_r + naccepted_t ; }

    /// get the number of rejected steps during the last call to generateEnsemble (after burn-in)
    int getNrejectedSteps() { return nrejected_r + nrejected_t ; }

    /// get the fraction of accepted steps during the last call to generateEnsemble (after burn-in)
    double getFractionAccepted() { 
      if( getNsteps() <= 0 ) return 0 ;
      return 1.0*getNacceptedSteps()/getNsteps() ; 
    }

    /// get the fraction of steps that were accepted when a radial step was performed during the last call to generateEnsemble (after burn-in)
    double getFractionAccepted_radial() { 
      if( naccepted_r + nrejected_r <= 0 ) return 0 ;
      return 1.0*naccepted_r/(naccepted_r+nrejected_r) ; 
    }

    /// get the fraction of steps that were accepted when a tangential step was performed during the last call to generateEnsemble (after burn-in)
    double getFractionAccepted_tangential() { 
      if( naccepted_t+nrejected_t <= 0 ) return 0 ;
      return 1.0*naccepted_t/(naccepted_t+nrejected_t) ; 
    }

    /// make one Markov chain step for a path
    bool doMarkovStep( JSourceModel* src, JScatteringModel* sm, JTargetModel* trg, double lambda_abs, JPhotonPath& p ) ;

    virtual int randomPathChange(  JPhotonPath& p, JTargetModel* trg ) ;

    /// set the average step size in degrees for the impact point on the target
    void setTargetStepSize_deg( double val ) {
      stepsize_angle_target = val * M_PI / 180 ;
    }

    /// set the average step size in [m] in the radial direction for the scattering vertices
    void setRadialStepSize_m( double val ) {
      stepsize_r = val ;
    }

    /**
       Set the average step size theta in degrees for steps in the tangential direction 
       for the scattering vertices.
       When r * theta is larger than the radial step size, the effective theta will be
       temporarily set to a smaller value to avoid taking steps that are too large.
    **/
    void setTangentialStepSize_deg( double val ) {
      stepsize_angle = val * M_PI / 180 ;
    }

    /// activate or deactive coordinate remapping
    void setCoordinateRemapping( bool val=true ) {
      apply_coordinate_remapping = val ;
    }

    /// returns true when coordinate remapping is activated, false otherwise
    bool getCoordinateRemapping() { return apply_coordinate_remapping ; }


  private: 

    /// return R0; the length scale used in the coordinate remapping
    inline double getR0() { return 3 * stepsize_r ; }

    /**
       return the distance between the remapped vertex and the previous vertex
       given the distance between the not-remapped vertex and the previous vertex
    **/
    double getRemappedDistance(double r) {
      if( 3.0*sqrt(3)*r < getR0() ) return pow(r*getR0()*getR0(),1.0/3.0) ;
      return r + 2.0 / (3.0*sqrt(3.0))*getR0() ;
    }

    /// inverse of getRemappedDistance
    double getUnmappedDistance(double rprime) {
      if( sqrt(3.0)*rprime < getR0() ) return rprime*rprime*rprime/(getR0()*getR0()) ;
      return rprime - 2.0/(3.0*sqrt(3.0))*getR0() ;
    }

    /**
       Return the remapped vertex position of x. 
       xprev is the (not remapped!) coordinate of the previous vertex in the photon path.
    **/
    JPosition3D getRemappedPosition( const JPosition3D& xprev, const JPosition3D& x ) {
      double r = x.getDistance(xprev) ;
      double rprime = getRemappedDistance(r) ;
      return xprev + (rprime/r)*(x-xprev) ;
    }

    /**
       Inverse of getRemappedPosition. Again, xprev is the (not remapped!) coordinate of the previous 
       vertex in the photon path.
    **/
    JPosition3D getUnmappedPosition( const JPosition3D& xprev, const JPosition3D& xprime ) {
      double rprime = xprime.getDistance(xprev) ;
      double r = getUnmappedDistance(rprime) ;
      return xprev + (r/rprime)*(xprime-xprev) ;
    }

    /// returns a remapped version of the photon path
    JPhotonPath getRemappedPhotonPath(const JPhotonPath& p ) {
      JPhotonPath pprime(p) ;
      for( int i=1; i<(int)pprime.size()-1; ++i ) {
	pprime[i] = getRemappedPosition(p[i-1],p[i]) ;
      }
      return pprime ;
    }

    /// inverse of getRemappedPhotonPath
    JPhotonPath getUnmappedPhotonPath(const JPhotonPath& pprime ) {
      JPhotonPath p(pprime) ;
      for( int i=1; i<(int)p.size()-1; ++i ) {
	p[i] = getUnmappedPosition(p[i-1],p[i]) ;
      }
      return p ;
    }

    /**
       Returns the conversion factor J needed to compute the path probability density
       in the remapped coordinates, for a given vertex.

       xprime is the remapped vertex position
       xprev is the (not remapped!) position of the previous vertex

       In casual notation (the ' indicates remapped coordinates):

       rho(x) dV =  rho(x) r^2 dr dOmega
       =  rho(x) r^2 / r'^2  dr/dr'  r'^2 dr' dOmega
       =  rho( x(x') ) r^2 / r'^2 dr/dr' dV'          <--- since dOmega == dOmega'
       := J x rho( x(x') )  dV'
       := rho'(x') dV'
       Where J :=  r^2 / r'^2 dr/dr'

       The point of this is that we typically pick r proportional to r'^3, so that 
       dr/dr' is proportional to r'^2, which cancels the factor 1/r'^2 in rho'. 
       The new factor r^2 then cancels the typical 1/r^2 singularity present in rho.
    **/
    double getRemappingCorrection( const JPosition3D& xprev, const JPosition3D& xprime ) {
      double rprime = xprev.getDistance(xprime) ;
      double r = getUnmappedDistance(rprime) ;
      if( sqrt(3.0)*rprime >= getR0() ) {
	// J = r^2 / r'^2 * dr/dr'
	// dr/dr' = 1
	return r*r/(rprime*rprime) ;
      }
      // J = r^2 / r'^2 * dr/dr'
      // dr/dr' = 3 * r'^2/R0^2
      // ==> J = 3 * r^2 / R0^2 
      return 3*r*r/(getR0()*getR0()) ;
    }

    /**
       Return the remapping correction for an entire photon path (product of the
       remapping correction for the individual scattering vertices).

       p is the photon path in normal coordinates
       pprime is the remapped version of p
    **/
    double getRemappingCorrection( const JPhotonPath& p, const JPhotonPath& pprime ) {
      double J = 1 ;
      for( int i=1; i<(int)p.size()-1; ++i ) {
	J *= getRemappingCorrection(p[i-1],pprime[i]) ;
      }
      return J ;
    }

  
    bool apply_coordinate_remapping ;
    int naccepted_r ; // number of accepted radial steps
    int nrejected_r ; // number of rejected radial steps
    int naccepted_t ; // number of accepted tangential steps
    int nrejected_t ; // number of rejected tangential steps
    double stepsize_angle_target ; 
    double stepsize_r ;
    double stepsize_angle ;
  } ;

  int JMarkovPathGenerator::randomPathChange( JPhotonPath& p, JTargetModel* trg ) {
    int nscat = p.n-2 ;
    int type = 0 ;

    if( nscat>0 ) {
      // pick vertex from 1 to nscat+1 (so excluding the start and end point)
      int nv = gRandom->Integer(nscat)+1 ;

      // distance w.r.t. previous vertex
      double r = p[nv-1].getDistance(p[nv]) ;

      if( gRandom->Uniform() > 0.5 ) {
	// with P=1/2, take a step in the radial direction
	type = 1 ;

	// to avoid nasty phase space Jacobian issues, I use this rather primitive method
	// where I simply generate a new position with a step in a random direction, then
	// only apply the radial part of the change
	//
	// NOTE: simply choosing 
	// *) r+dr with P proportional to  (r+dr)^2 
	// *) r-dr with P proportional to  (r-dr)^2
	// does not work. Trust me. I've tried.
	double _r = gRandom->Exp(stepsize_r) ;
	double _x, _y, _z ;
	gRandom->Sphere( _x, _y, _z, _r ) ;
	double _rnew = ( JPosition3D(_x,_y,_z) + JPosition3D(0,0,r) ).getLength() ;
	double dr = _rnew - r ;

	p[nv] = p[nv] + dr * JPosition3D(JVersor3D(p[nv]-p[nv-1])) ;

      } else {
	// with P=1/2, take a step in the tangential direction
	type = 2 ;

	// for large radii, limit the maximum angle over which we can move
	// so that the average change in distance does not exceed stepsize_r too much
	double _stepsize_angle = min(stepsize_angle,stepsize_r/r) ;
	double theta = gRandom->Exp(_stepsize_angle) ;
	double phi   = gRandom->Uniform(2.0*M_PI) ;
	JVersor3D dir( sin(theta)*cos(phi), sin(theta)*sin(phi), cos(theta) ) ;
	JVersor3D xdir( p[nv] - p[nv-1] ) ;
	JRotation3D rot(xdir) ;
	dir = dir.rotate_back(rot) ;
	p[nv] = p[nv-1] + r * JVector3D(dir) ;
      }
    }

    // in case of a finite target, also move the end point
    if( trg->getRadius() > 0 ) {
      const double dct = 1-cos(stepsize_angle_target) ;
      double ct = -2 ;
      while( ct<=-1 ) ct = 1-gRandom->Exp(dct) ;
      double theta = acos(ct) ;
      double phi = gRandom->Uniform(2*M_PI) ;
      JVersor3D original_dir( p.back()-trg->getPosition() ) ;
      const JRotation3D R( original_dir ) ;
      JVersor3D new_dir( JAngle3D(theta,phi) ) ;
      new_dir = new_dir.rotate_back(R) ;
      p.back() = trg->getPosition() + trg->getRadius() * JVector3D(new_dir) ;
    }

    return type ;
  }

  std::vector<JPhotonPath> JMarkovPathGenerator::generateEnsemble(  int n, const JPhotonPath& start_path, JSourceModel* src, JScatteringModel* sm, JTargetModel* trg, double lambda_abs, int nsteps_burn_in, int nsteps_save ) {
    vector<JPhotonPath> result ;

    JPhotonPath testpath(start_path) ;

    // check that the starting path is okay
    double dist = testpath[0].getDistanceSquared(src->getPosition()) ;
    if( dist != 0 ) {
      cerr << "ERROR in generateEnsemble: testpath position " << testpath[0] 
	   << "does not match source position " << src->getPosition() << "!" << endl ;
      return result ;
    }
    {
      double rho = getPhotonPathProbabilityDensity(testpath,src,sm,trg,lambda_abs) ;
      if( rho  == 0 ) {
	cerr << "ERROR in generateEnsemble: testpath probability density = 0" << endl ;
	return result ;
      }
    }

    // convert the path to remapped coordinates
    if( apply_coordinate_remapping ) testpath = getRemappedPhotonPath(testpath) ;

    // burn-in
    int i = 0 ;
    while( i!=nsteps_burn_in ) {
      doMarkovStep(src,sm,trg,lambda_abs,testpath) ; 
      ++i ;
    }

    // ensemble generation
    i = 0 ;
    naccepted_r = 0 ;
    nrejected_r = 0 ;
    naccepted_t = 0 ;
    nrejected_t = 0 ;
    while( i!=n ) {
      // take nsteps_save steps
      int j = 0 ;
      while( j!=nsteps_save ) {
	if( doMarkovStep(src,sm,trg,lambda_abs,testpath) ) {
	  ++j ;
	}
      }
      // save the path to the ensemble
      if( apply_coordinate_remapping ) {
	result.push_back( getUnmappedPhotonPath(testpath) ) ;
      } else {
	result.push_back( testpath ) ;
      }
      ++i ;
    }
    
    return result ;
  }

  bool JMarkovPathGenerator::doMarkovStep( JSourceModel* src, JScatteringModel* sm, JTargetModel* trg, double lambda_abs, JPhotonPath& p ) {
    int nscat = p.n-2 ;
    if( nscat == 0 && !trg->getRadius()>0 ) {
      return false ;
    }

    // copy path
    JPhotonPath copy = p ;

    double rhoBefore ;
    if( !apply_coordinate_remapping ) {
      rhoBefore = getPhotonPathProbabilityDensity(copy,src,sm,trg,lambda_abs) ;
    } else {
      // we are working in remapped coordinates, so we have to calculate a corrected
      // probability density
      JPhotonPath copy_unmapped = getUnmappedPhotonPath(copy) ;
      rhoBefore = getPhotonPathProbabilityDensity(copy_unmapped,src,sm,trg,lambda_abs) ;
      rhoBefore *= getRemappingCorrection(copy_unmapped,copy) ;
    }

    if( rhoBefore == 0 ) {
      cerr << "FATAL ERROR: starting probability density = 0" << endl ;
      exit(1) ;
    }

    // apply random change to the path
    int type = randomPathChange(copy,trg) ;

    // recalculate log L
    double rhoAfter ;
    if( !apply_coordinate_remapping ) {
      rhoAfter = getPhotonPathProbabilityDensity(copy,src,sm,trg,lambda_abs) ;
    } else {
      // we are working in remapped coordinates, so we have to calculate a corrected
      // probability density
      JPhotonPath copy_unmapped = getUnmappedPhotonPath(copy) ;
      rhoAfter = getPhotonPathProbabilityDensity(copy_unmapped,src,sm,trg,lambda_abs) ;
      rhoAfter *= getRemappingCorrection(copy_unmapped,copy) ;
    }

    // use Metropolis-Hastings algorithm to keep or reject change
    if( rhoAfter > rhoBefore ) {
      // accept the change
      p = copy ;
      if( type <  2 ) ++naccepted_r ;
      if( type == 2 ) ++naccepted_t ;
      return true ;
    } else {
      // accept the change with some probability
      double P = rhoAfter/rhoBefore ;
      if( gRandom->Uniform()<P ) {
	p = copy ;
	if( type <  2 ) ++naccepted_r ;
	if( type == 2 ) ++naccepted_t ;
	return true ;
      }
    }
    if( type <  2 ) ++nrejected_r ;
    if( type == 2 ) ++nrejected_t ;
    return false ;
  }

}

#endif