#ifndef H_MARKOV_GENERATOR
#define H_MARKOV_GENERATOR

/**
 * \author mjongen
 */

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

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

// JPP includes
#include "JGeometry3D/JPosition3D.hh"

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

namespace JMARKOV {

  using namespace JGEOMETRY3D ;

  /**
     Abstract interface for the generation of points in 3D space.
     By convention, the random number generator used in the derived
     classes is gRandom.
   **/
  class JGenerator {
    public:

    /**
       Return a randomly generated position
    **/
    virtual JPosition3D getPosition() = 0 ;

    /**
       return the weight (=probability density dP/dV) for the
       given position. 
       When the generator is properly normalized, the integral
       of this quantity over the whole space is 1.
     **/
    virtual double getWeight( JPosition3D pos ) = 0 ;

  protected :
    
  } ;

  /**
     Implementation of the JGenerator interface.

     Generates positions uniformly within a rectangular cuboid,
   **/
  class JUniformGenerator : public JGenerator {

    public :

    /**
       Constructor. The arguments are the limits of each coordinate.
     **/
    JUniformGenerator( double xmin, double ymin, double zmin, double xmax, double ymax, double zmax ) : posmin(xmin,ymin,zmin), posmax(xmax,ymax,zmax) {
      setVolume() ;
    }

    /**
       Constructor. The arguments are the 'lower left bottom' corner and the 
       'upper right top' corner.
     **/
    JUniformGenerator( JPosition3D _posmin, JPosition3D _posmax ) : posmin(_posmin), posmax(_posmax) {
      setVolume() ;
    }

    JPosition3D getPosition() {
      double x = gRandom->Uniform( posmin.getX(), posmax.getX() ) ;
      double y = gRandom->Uniform( posmin.getY(), posmax.getY() ) ;
      double z = gRandom->Uniform( posmin.getZ(), posmax.getZ() ) ;
      return JPosition3D(x,y,z) ;
    }

    double getWeight( JPosition3D pos ) {
      if( pos.getX() < posmin.getX() ||  pos.getX() > posmax.getX() ) return 0.0 ;
      if( pos.getY() < posmin.getY() ||  pos.getY() > posmax.getY() ) return 0.0 ;
      if( pos.getZ() < posmin.getZ() ||  pos.getZ() > posmax.getZ() ) return 0.0 ;
      return 1.0/V ;
    }

  protected :

    void setVolume() {
      V = (posmax.getX()-posmin.getX())*(posmax.getY()-posmin.getY())*(posmax.getZ()-posmin.getZ()) ;
    }

    JPosition3D posmin ;
    JPosition3D posmax ;
    double V ; // the total volume
  } ;

  /**
     Implementation of the JGenerator interface.

     Generates positions uniformly within a ball of a given radius
     centered around the origin.
   **/
  class JBallGenerator : public JGenerator {
    public :
    
    /**
       Constructor. The argument is the radius of the ball.
     **/
    JBallGenerator( double _R ) : R(_R), Rcube(_R*_R*_R) {
      setVolume() ;
    }

    JPosition3D getPosition() {
      double rcube = gRandom->Uniform(0,Rcube) ;
      double r = pow(rcube,1.0/3.0) ;
      double x,y,z ;
      gRandom->Sphere(x,y,z,r) ;
      return JPosition3D(x,y,z) ;
    }

    double getWeight( JPosition3D pos ) {
      if( pos.getLength()>R ) return 0.0 ;
      return 1.0/V ;
    }

    protected :

    void setVolume() { V = 4.0/3.0*M_PI*Rcube ; }

    const double R ;
    const double Rcube ;
    double V ;
  } ;

  /**
     Implementation of the JGenerator interface.

     Generates positions according to the distribution

     1/r^2 exp(-r/lambda) 

     around the origin.
   **/
  class JExpRsqInvGenerator : public JGenerator {
    public :
    
    /**
       Constructor. The argument is the radius of the ball.
     **/
    JExpRsqInvGenerator( double _lambda ) : lambda(_lambda) {}

    JPosition3D getPosition() {
      double xi = gRandom->Uniform() ;
      double r = -lambda*log(1-xi) ;
      double x,y,z ;
      gRandom->Sphere(x,y,z,r) ;
      return JPosition3D(x,y,z) ;
    }

    double getWeight( JPosition3D pos ) {
      double r = pos.getLength() ;
      return lambda*exp(-r/lambda)/(r*r*4*M_PI) ;
    }

    protected :

    const double lambda ;
  } ;

  /**
     The 'magical distributions' are a class of distributions.

     The magical distribution of degree N is the distribution of the distance
     to the origin after taking N random steps where the step size r is drawn
     from an "exponential inverse r-squared"  1/r^2 exp(-r/lambda) 
     distribution, which is the base behaviour for scattering 

     Note that it seems to satisfy the property (for lambda = 1 )
     mean = -0.5 + sqrt( (1+8*n)/4 )
   **/
  class JMagicalDistribution : public JGenerator {

   public :

    JMagicalDistribution( unsigned int _N, double _lambda ) : N(_N), lambda(_lambda), gen(1) {
      const int nbins = 1000 ;
      const double xmax = 100 ;
      const int nsamples = 10000000 ;
      h = new TH1F("h","",nbins,0,xmax) ;
      for( int i=0 ; i<nsamples ; ++i ) {
	JPosition3D pos = gen.getPosition() ;
	for( unsigned int n=1 ; n<N ; ++n ) pos += gen.getPosition() ;
	double r = pos.getLength() ;
	h->Fill(r) ;
      }
      h->Scale( 1.0/h->Integral("width") ) ;
    }

    ~JMagicalDistribution() {
      delete h ;
    }

    JPosition3D getPosition() {
      double r = h->GetRandom() ;
      double x,y,z ;
      gRandom->Sphere(x,y,z,r) ;
      return JPosition3D(x,y,z) ;
    }

    double getWeight( JPosition3D pos ) {
      double r =  pos.getLength() ;
      Int_t bin = h->FindBin(r) ;
      return h->GetBinContent(bin)/(r*r) ;
    }

  protected :

    unsigned int N ; // degree
    double lambda ;
    JExpRsqInvGenerator gen ;
    TH1F* h ;
  } ;

  /**
     Implementation of the JGenerator interface.

     Generates positions according to a probability density proportional
     to 1/r^2 centered around some point.
   **/
  class JSingularityGenerator : public JGenerator {
    public :
    
    /**
       Constructor. The argument is the maximal distance from _x0
     **/
    JSingularityGenerator( double _R, JPosition3D _x0 ) : R(_R), C(1.0/(4*M_PI*R)), x0(_x0) {}

    JPosition3D getPosition() {
      double r = gRandom->Uniform(0,R) ;
      double x,y,z ;
      gRandom->Sphere(x,y,z,r) ;
      return JPosition3D(x,y,z)+x0 ;
    }

    double getWeight( JPosition3D pos ) {
      pos -= x0 ;
      if( pos.getLength()>R ) return 0.0 ;
      return C/pos.getLengthSquared() ;
    }

    protected :

    const double R ;
    const double C ; // constant factor in weight calculation
    JPosition3D x0 ;
  } ;

  /**
     Implementation of the JGenerator interface.

     Generates positions according to an exponential distribution,
     around the origin, i.e. with probability density

     rho( vec x ) = alpha * exp(-r/l)

     for given l. The minimal and maximal values of r are given by
     r1 and r2, respectively.
   **/
    class JExponentialGenerator : public JGenerator {

    public :

    /**
       Constructor. 
       _r1 and _r2 are the minimal and maximal allowed value of r.
       _L is the length scale.
    **/
    JExponentialGenerator( double _r1, double _r2, double _L ) :
      r1(_r1), r2(_r2), L(_L) 
    {
      // normalization factor
      ximin = getIntegrand(r1) ;
      ximax = getIntegrand(r2) ;
      C = ximax-ximin ;
    }

    double getWeight( JPosition3D pos ) {
      return getWeight( pos.getLength() ) ;
    }

    JPosition3D getPosition() {
      double x, y, z ;
      double r = getR() ;
      gRandom->Sphere(x,y,z,r) ;
      return JPosition3D(x,y,z) ;
    }

    double getR() {
      // generate integrand value
      double xi = gRandom->Uniform(ximin,ximax) ;
      // find matching value of r
      return getInvertedIntegrand(xi) ;
    }

    double getWeight( double r ) {
      return exp(-r/L) / C ;
    }


  protected :

    double getIntegrand( double r ) {
      return -4*M_PI*L*exp(-r/L)*(2*L*L+2*L*r+r*r) ;
    }

    /// return value y such that getIntegrand(y) = x
    double getInvertedIntegrand( double x ) {
      const double precision = 1e-10 ;
      // using the bisection algorithm
      double rl = r1 ;
      double vl = getIntegrand(rl)-x ;
      double rr = r2 ;
      double vr = getIntegrand(rr)-x ;
      double rc ;
      double vc ;

      while( rr-rl > precision ) {
	rc = 0.5*(rl+rr) ;
	vc = getIntegrand(rc)-x ;
	if( (vc>0) == (vr>0) ) {
	  // must be in left half
	  rr = rc ;
	  vr = vc ;
	} else {
	  // must be in right half
	  rl = rc ;
	  vl = vc ;
	}
      }
      return rc ;
    }

    double r1 ;
    double r2 ;
    double L ;
    double C ;
    double ximin ;
    double ximax ;
  } ;

  /**
     Implementation of the JGenerator interface.

     Generates positions according to a 3D Gaussian distribution
     centered around the origin.
   **/
  class JGaussianGenerator : public JGenerator {
    public :
    
    /**
       Constructor. The argument is the width of the gaussian.
     **/
    JGaussianGenerator( double _sigma ) : sigma(_sigma) {
      C = 1.0 / ( 2*M_PI*sqrt(2*M_PI) * sigma * sigma * sigma ) ;
    }

    JPosition3D getPosition() {
      double x = gRandom->Gaus(0,sigma) ;
      double y = gRandom->Gaus(0,sigma) ;
      double z = gRandom->Gaus(0,sigma) ;
      return JPosition3D(x,y,z) ;
    }

    double getWeight( JPosition3D pos ) {
      return C*exp(-0.5*pos.getLengthSquared()/(sigma*sigma)) ;
    }

    protected :

    const double sigma ;
    double C ;
  } ;

  /**
     Implementation of the JGenerator interface.

     Generates positions by combining two other JGenerators
   **/
  class JCombinedGenerator : public JGenerator {
    public :
    
    /**
       Constructor.
       
       The arguments _gn are pointers to the two JGenerators to use.
       Note that ownership is not transferred, i.e. 
       - the memory taken by the _gn is not freed when JCombinedGenerator 
         is deleted
       - changes made to the instances pointed to by the _gn will affect
         the behaviour of the JCombinedGenerator

       The arguments _cn are the relative weights of the two generators.
       They do not have to add up to 1.
       The probability that a value is generated from g1 is 
       c1/(c1+c2).
     **/
    JCombinedGenerator( double _c1, JGenerator* _g1, double _c2, JGenerator* _g2 ) : g1(_g1), g2(_g2) {
      p1 = _c1/(_c1+_c2) ;
      p2 = 1-p1 ;
    }

    JPosition3D getPosition() {
      if( gRandom->Uniform()<p1 ) {
	return g1->getPosition() ;
      } else {
	return g2->getPosition() ;
      }
    }

    double getWeight( JPosition3D pos ) {
      double w1 = g1->getWeight(pos) ;
      double w2 = g2->getWeight(pos) ;
      return p1*w1 + p2*w2 ;
    }

    protected :

    JGenerator* g1 ;
    JGenerator* g2 ;
    double p1 ;
    double p2 ;
  } ;

  /**
     Implementation of the JGenerator interface.

     Generates positions by combining three other JGenerators
   **/
  class JTripleGenerator : public JGenerator {
    public :
    /**
       Constructor.
       
       The arguments gn are pointers to the three JGenerators to use.
       Note that ownership is not transferred, i.e. 
       - the memory taken by the _gn is not freed when JTripleGenerator 
         is deleted
       - changes made to the instances pointed to by the _gn will affect
         the behaviour of the JTripleGenerator

       The arguments cn are the relative weights of the two generators.
       They do not have to add up to 1.
       The probability that a value is generated from gi is 
       ci/(c1+c2+c3).
     **/
    JTripleGenerator( double c1, JGenerator* g1, double c2, JGenerator* g2, double c3, JGenerator* g3 ) : 
      gsub(c1,g1,c2,g2), g(c1+c2,&gsub,c3,g3) {}

    JPosition3D getPosition() { return g.getPosition() ; }

    double getWeight( JPosition3D pos ) { return g.getWeight(pos) ; }

    protected :

    JCombinedGenerator gsub ;
    JCombinedGenerator g ;
  } ;

  /**
     Implementation of the JGenerator interface.

     Generates positions by shifting the positions from another JGenerator
   **/
  class JShiftedGenerator : public JGenerator {
    public :
    
    /**
       Constructor.
       
       _g is a pointer to some other generator, note that
       memory ownership of the instance pointed to by _g 
       is not transferred, so
       - it will not be deleted with JShiftedGenerator
       - changes to _g affect JShiftedGenerator as well
     **/
    JShiftedGenerator( JGenerator* _g, JPosition3D _shift ) : 
      g(_g), shift(_shift) {}

    JPosition3D getPosition() {
      return g->getPosition() + shift ;
    }

    double getWeight( JPosition3D pos ) {
      return g->getWeight( pos-shift ) ;
    }

    protected :

    JGenerator* g ;
    JPosition3D shift ;
  } ;

  /**
     Implementation of the JGenerator interface.

     Generates positions by sampling from 3 histograms
     that represent the projection of the distribution along the
     X, Y and Z axis.
   **/
  class JHistGenerator : public JGenerator {
    public :
    
    /**
       Constructor. 
       The histograms are cloned and automatically normalized.

       Note: empty bins should always be avoided!
     **/
    JHistGenerator( TH1* _hx, TH1* _hy, TH1* _hz ) {
      // copy histograms
      hx = (TH1*) _hx->Clone( _hx->GetName() ) ;
      hy = (TH1*) _hy->Clone( _hy->GetName() ) ;
      hz = (TH1*) _hz->Clone( _hz->GetName() ) ;
      // normalize histograms
      hx->Scale( 1.0/hx->Integral("width") ) ;
      hy->Scale( 1.0/hy->Integral("width") ) ;
      hz->Scale( 1.0/hz->Integral("width") ) ;
    }

    ~JHistGenerator() {
      delete hx ;
      delete hy ;
      delete hz ;
    }

    JPosition3D getPosition() {
      double x = hx->GetRandom() ;
      double y = hy->GetRandom() ;
      double z = hz->GetRandom() ;
      return JPosition3D(x,y,z) ;
    }

    double getWeight( JPosition3D pos ) {
      Int_t bin ;
      double w = 1 ;
      
      bin = hx->GetXaxis()->FindBin( pos.getX() ) ;
      if( bin==0 || bin==hx->GetNbinsX()+1 ) return 0.0 ; // out of range
      w *= hx->GetBinContent(bin) ;

      bin = hy->GetXaxis()->FindBin( pos.getY() ) ;
      if( bin==0 || bin==hy->GetNbinsX()+1 ) return 0.0 ; // out of range
      w *= hy->GetBinContent(bin) ;

      bin = hz->GetXaxis()->FindBin( pos.getZ() ) ;
      if( bin==0 || bin==hz->GetNbinsX()+1 ) return 0.0 ; // out of range
      w *= hz->GetBinContent(bin) ;
 
      return w ;
    }

    TH1* hx ;
    TH1* hy ;
    TH1* hz ;
  } ;

  /**
     Implementation of the JGenerator interface.

     Generates positions on a sphere of radius _r centered at _x0 
     by sampling from a 2D histogram of cosTheta vs phi.

     When r=0, it simply returns the position x0.
   **/
  class JSphereGenerator : public JGenerator {
    public :
    
    /**
       Constructor. 
       The histogram is cloned and automatically normalized.
       The x-axis is assumed to be cos(theta), and the y-axis is 
       assumed to be phi in radians.
       Empty bins should be avoided!
     **/
    JSphereGenerator( const JPosition3D& _x0, double _r=0, TH2* _h=NULL ) : x0(_x0), r(_r) {
      h = NULL ;

      if( r>0 ) {
	// copy histogram
	h = (TH2*) _h->Clone( _h->GetName() ) ;
	// check that it has the proper limits
	if( h->GetXaxis()->GetXmin()<-1 || h->GetXaxis()->GetXmax()>1 ) {
	  cerr << "FATAL ERROR in JSphereGenerator. Invalid x-axis range." << endl ;
	  exit(1) ;
	}
	if( h->GetYaxis()->GetXmin() < -M_PI || h->GetYaxis()->GetXmax() > M_PI ) {
	  cerr << "FATAL ERROR in JSphereGenerator. Invalid y-axis range." << endl ;
	  exit(1) ;
	}
	// normalize histogram
	h->Scale( 1.0/h->Integral("width") ) ;
	h->SetOption("colz") ;
      }
    }

    ~JSphereGenerator() {
      delete h ;
    }

    JPosition3D getPosition() {
      if( r==0 ) return x0 ;

      double costheta, phi ;
      h->GetRandom2(costheta,phi) ;
      return x0 + r * JVector3D(JVersor3D(JAngle3D(acos(costheta),phi))) ;
    }

    double getWeight( JPosition3D pos ) {
      if( r==0 ) return 1 ;

      // check that the position is in fact on the target surface
      const double tolerance = 1e-5 ;
      if( fabs( pos.getDistance(x0)-r ) > tolerance*r ) {
	return 0.0 ;
      }

      JVersor3D dir(pos-x0) ;
      double ct = dir.getDZ() ;
      double phi = dir.getPhi() ;

      if( ct < h->GetXaxis()->GetXmin() || ct >=h->GetXaxis()->GetXmax() ) return 0.0 ;
      if( phi< h->GetYaxis()->GetXmin() || phi>=h->GetYaxis()->GetXmax() ) return 0.0 ;

      Int_t bin = h->FindBin(ct,phi) ;
      return h->GetBinContent(bin) / (r*r) ;
    }

    JPosition3D x0 ;
    double r ;
    TH2* h ;
  } ;

  /**
     Implementation of the JGenerator interface.

     Generates positions uniformly by sampling from a histogram
     that represents the x,y,z distribution
   **/
  class J3DhistGenerator : public JGenerator {
    public :
    
    /**
       Constructor. 
       The histogram is cloned and automatically normalized.

       Note: empty bins should always be avoided!
     **/
    J3DhistGenerator( TH3* _h ) {
      // copy histogram
      h = (TH3*) _h->Clone( _h->GetName() ) ;
      // normalize histogram
      h->Scale( 1.0/h->Integral("width") ) ;
    }

    ~J3DhistGenerator() {
      delete h ;
    }

    JPosition3D getPosition() {
      double x, y, z ;
      h->GetRandom3(x,y,z) ;
      return JPosition3D(x,y,z) ;
    }

    double getWeight( JPosition3D pos ) {
      // return 0 outside histogram range
      if( pos.getX()<h->GetXaxis()->GetXmin() || pos.getX()>=h->GetXaxis()->GetXmax() ) return 0.0 ;
      if( pos.getY()<h->GetYaxis()->GetXmin() || pos.getY()>=h->GetYaxis()->GetXmax() ) return 0.0 ;
      if( pos.getZ()<h->GetZaxis()->GetXmin() || pos.getZ()>=h->GetZaxis()->GetXmax() ) return 0.0 ;
      // return weight
      Int_t bin = h->FindBin( pos.getX(), pos.getY(), pos.getZ() ) ;
      return h->GetBinContent(bin) ;
    }

    TH3* h ;
  } ;

}

#endif