#ifndef STOCHASTIC_RUNGE_KUTTA
#define STOCHASTIC_RUNGE_KUTTA

#include <vector>
#include <cmath>
#include "InternalDefs.h"
#include "Auxiliary.cpp"
using namespace std;





//Two step, explicit, stochastic Runge-Kutta integration with fixed time step, with mean-square order of accuracy 3/2.
//See [Milstein 1995, Numerical Integration of Stochastic Differential Equations, Eq. (3.55) in Theorem 3.3]
//Returns false on error
//VECTORFIELD_FUNCTOR should implement: 
//	bool operator()(double time, const POINT &p, POINT &velocity, POINT &finalPoint) const,
//  double D(unsigned int) const,
//where D(k) returns the time-independent diffusion coefficient (white noise amplitude = SQRT(2D)) for component k.
//finalPoint is a pure return value. If the functor returns false, finalPoint will hold a suggested "last valid point" at which the trajectory is suspected to have left the domain of the vector field.
//
//POINT should implement:
//	POINT operator+(POINT, POINT),
//	POINT operator-(POINT, POINT),
//	POINT operator*(double,POINT),
//satisfying the axioms of a real vector field. Additionally, POINT should implement,
//  double& operator[](unsigned int k)

template<class POINT, class VECTORFIELD_FUNCTOR>
bool StochasticRungeKutta(	const POINT 				&startPoint,
							double 						maxT, //max integration time
							double 						dt, //fixed time step
							const VECTORFIELD_FUNCTOR 	&f,
							unsigned int 				dimension, //finite system dimension, i.e. number of components, of POINT
							bool 						onlySaveFinalPoint,
							unsigned long				maxRecordedPoints, //if small, some intermediate trajectory points are skipped (i.e. not recorded). This might be useful if accurate integration requires a small time step, but would produce to large of a time series.
							vector<double> 				&realizedTimes, //if (onlySaveFinalPoint==false), realized time-points, corresponding to realizedPath[]. If (onlySaveFinalPoint==true), will only contain last time.
							vector<POINT> 				&realizedPath, //if (onlySaveFinalPoint==false), traversed points after integration, corresponding to realizedTimes[]. If (onlySaveFinalPoint==true), will only contain last path point.
							bool						&reachedDomainBoundary // (output) returns true if the domain of f() was exited during integration. If that's the case, realizedTimes[] and realizedPath[] will contain that last point.
							){
		
	//preliminary error checking
	if(dt<EPSILON) return false;
	if(maxT < dt) return false;
	if(dimension<1) return false;
	if(maxRecordedPoints < 1) return false;
	
	POINT currentPoint, tempPoint1, tempPoint2, finalPoint;
	POINT k0,k1,k2,k3;
	double t;
	bool lastUnsaved;
	double *sigmas 			= new double[dimension];
	double *xis				= new double[dimension];
	double *etas			= new double[dimension];
	unsigned int n;
	unsigned long steps;
	
	unsigned long skipStep = std::max((unsigned long)1, (unsigned long)ceil((maxT/dt)/maxRecordedPoints));
	
	
	//initialization
	for(n=0; n<dimension; ++n){
		sigmas[n] 	= std::sqrt(2.0*f.D(n));
	}
	realizedTimes.clear();
	realizedPath.clear();
	reachedDomainBoundary = false;
	lastUnsaved = false;
	currentPoint = startPoint;
	if(!onlySaveFinalPoint){
		realizedTimes.push_back(0);
		realizedPath.push_back(startPoint);
	}
	
	//run
	for(steps=1, t = 0; t<maxT - EPSILON; t+=dt, ++steps){
		//generate random numbers
		for(n=0; n<dimension; ++n){
			xis[n] 	= standardNormal() * sigmas[n] * std::sqrt(dt);
			etas[n]	= standardNormal() * sigmas[n] * std::sqrt(dt/12.0);
		}
		
		//iterate
		if(!f(t, currentPoint, k0, finalPoint)){
			reachedDomainBoundary = true;
			goto END_StochasticRungeKutta2;
		}
		
		lastUnsaved = true;
		tempPoint1 = currentPoint + dt*k0;
		for(n=0; n<dimension; ++n){ tempPoint1[n] += xis[n]; }
		if(!f(t, tempPoint1, k1, finalPoint)){
			reachedDomainBoundary = true;
			currentPoint = tempPoint1;
			goto END_StochasticRungeKutta2;
		}
		
		tempPoint2 = currentPoint;
		for(n=0; n<dimension; ++n){
			tempPoint2[n] += sigmas[n]*dt;
		}
		if(!f(t, tempPoint2, k2, finalPoint)){
			reachedDomainBoundary = true;
			currentPoint = tempPoint2;
			goto END_StochasticRungeKutta2;
		}
		
		for(n=0; n<dimension; ++n){
			k3 = (n==0 ? etas[n]*(k2-k0) : k3 + etas[n]*(k2-k0));
		}
		currentPoint = currentPoint + (dt/2)*(k0 + k1) + k3;
		for(n=0; n<dimension; ++n){ currentPoint[n] += xis[n]; }
		
		//possibly save current point
		if((!onlySaveFinalPoint) && (steps % skipStep == 0)){
			realizedTimes.push_back(t+dt);
			realizedPath.push_back(currentPoint);
			lastUnsaved = false;
		}
	}
	
	END_StochasticRungeKutta2:
	
	if(lastUnsaved){
		realizedTimes.push_back(t);
		realizedPath.push_back(reachedDomainBoundary ? finalPoint : currentPoint);
	}
	//do one last domain boundary check for final point
	if(!reachedDomainBoundary){			
		if(!f(t, currentPoint, k0, finalPoint)){
			reachedDomainBoundary = true;
			realizedPath[realizedPath.size()-1] = finalPoint;
		}
	}
	
	delete [] sigmas;
	delete [] xis;
	delete [] etas;
	
	return true;
}










//Same as StochasticRungeKutta(..) above, but with variable noise, i.e. here white noise amplitude is allowed to depend on the current state
//Hence, the object f needs to implement    f.D(unsigned int k, const POINT &point) const
template<class POINT, class VECTORFIELD_FUNCTOR>
bool StochasticRungeKutta_VN(	const POINT 				&startPoint,
								double 						maxT, //max integration time
								double 						dt, //fixed time step
								const VECTORFIELD_FUNCTOR 	&f,
								unsigned int 				dimension, //finite system dimension, i.e. number of components, of POINT
								bool 						onlySaveFinalPoint,
								unsigned long				maxRecordedPoints, //if small, some intermediate trajectory points are skipped (i.e. not recorded). This might be useful if accurate integration requires a small time step, but would produce to large of a time series.
								vector<double> 				&realizedTimes, //if (onlySaveFinalPoint==false), realized time-points, corresponding to realizedPath[]. If (onlySaveFinalPoint==true), will only contain last time.
								vector<POINT> 				&realizedPath, //if (onlySaveFinalPoint==false), traversed points after integration, corresponding to realizedTimes[]. If (onlySaveFinalPoint==true), will only contain last path point.
								bool						&reachedDomainBoundary // (output) returns true if the domain of f() was exited during integration. If that's the case, realizedTimes[] and realizedPath[] will contain that last point.
								){
		
	//preliminary error checking
	if(dt<EPSILON) return false;
	if(maxT < dt) return false;
	if(dimension<1) return false;
	if(maxRecordedPoints < 1) return false;
	
	POINT currentPoint, tempPoint1, tempPoint2, finalPoint;
	POINT k0,k1,k2,k3;
	double t;
	bool lastUnsaved;
	double *sigmas 			= new double[dimension];
	double *xis				= new double[dimension];
	double *etas			= new double[dimension];
	unsigned int n;
	unsigned long steps;
	
	unsigned long skipStep = std::max((unsigned long)1, (unsigned long)ceil((maxT/dt)/maxRecordedPoints));
	
	realizedTimes.clear();
	realizedPath.clear();
	reachedDomainBoundary = false;
	lastUnsaved = false;
	currentPoint = startPoint;
	if(!onlySaveFinalPoint){
		realizedTimes.push_back(0);
		realizedPath.push_back(startPoint);
	}
	
	//run
	for(steps=1, t = 0; t<maxT - EPSILON; t+=dt, ++steps){
		//get white noise amplitudes
		for(n=0; n<dimension; ++n){
			sigmas[n] 	= std::sqrt(2.0*f.D(n,currentPoint));
		}
	
		//generate random numbers
		for(n=0; n<dimension; ++n){
			xis[n] 	= standardNormal() * sigmas[n] * std::sqrt(dt);
			etas[n]	= standardNormal() * sigmas[n] * std::sqrt(dt/12.0);
		}
		
		//iterate
		if(!f(t, currentPoint, k0, finalPoint)){
			reachedDomainBoundary = true;
			goto END_StochasticRungeKutta2;
		}
		
		lastUnsaved = true;
		tempPoint1 = currentPoint + dt*k0;
		for(n=0; n<dimension; ++n){ tempPoint1[n] += xis[n]; }
		if(!f(t, tempPoint1, k1, finalPoint)){
			reachedDomainBoundary = true;
			currentPoint = tempPoint1;
			goto END_StochasticRungeKutta2;
		}
		
		tempPoint2 = currentPoint;
		for(n=0; n<dimension; ++n){
			tempPoint2[n] += sigmas[n]*dt;
		}
		if(!f(t, tempPoint2, k2, finalPoint)){
			reachedDomainBoundary = true;
			currentPoint = tempPoint2;
			goto END_StochasticRungeKutta2;
		}
		
		for(n=0; n<dimension; ++n){
			k3 = (n==0 ? etas[n]*(k2-k0) : k3 + etas[n]*(k2-k0));
		}
		currentPoint = currentPoint + (dt/2)*(k0 + k1) + k3;
		for(n=0; n<dimension; ++n){ currentPoint[n] += xis[n]; }
		
		//possibly save current point
		if((!onlySaveFinalPoint) && (steps % skipStep == 0)){
			realizedTimes.push_back(t+dt);
			realizedPath.push_back(currentPoint);
			lastUnsaved = false;
		}
	}
	
	END_StochasticRungeKutta2:
	
	if(lastUnsaved){
		realizedTimes.push_back(t);
		realizedPath.push_back(reachedDomainBoundary ? finalPoint : currentPoint);
	}
	//do one last domain boundary check for final point
	if(!reachedDomainBoundary){			
		if(!f(t, currentPoint, k0, finalPoint)){
			reachedDomainBoundary = true;
			realizedPath[realizedPath.size()-1] = finalPoint;
		}
	}
	
	delete [] sigmas;
	delete [] xis;
	delete [] etas;
	
	return true;
}



#endif