#  Code for Stochastic Calculus class, Jonathan Goodman
#  http://www.math.nyu.edu/faculty/goodman/teaching/StochCalc2022/
#  PathDemo.py
#  uses Python 3 and Numpy 
#  For Section 1

#  Demo of paths and probability densities for dX = -aXdt + sigma*dW

# illustrate using the Python random number generator from numpy version 19

from   numpy.random import default_rng         #  numpy randon number generator routines
import numpy             as np
import matplotlib.pyplot as plt



#-------------------- The simulator, Create OU paths using EM  --------------------

def path_sim( X0, n, dt, T, a, sig, rng):
    """Simulate and return n independent paths for the Ornstein Unhelbeck
       process dX_t = -a X_t dt + sig*dW_t.
    Arguments
    X0:  (type = float) start value of the paths, X[i,0] = X0
    n:   (type = int) the number of paths
    dt:  (type = float) time step for the simulator
    a:   (type = float) mean reversion parameter
    sig: (type = float) noise parameter
    rng: an instance of a random number generator: np.random.Generator
    return: (X,T) (type = tuple, (np array, np array))
            X = simuated paths, X[i,k] = step k of path i
            T = times T[k] = time of step k
      Note: X[:,0] = X0[:] (copy starting points) and T[0] = 0
    """
    
#      adjust dt (make it slightly smaller) to get an integer number of steps

    nT = T/dt        #  "number" of time steps, but may not be an integer
    nT = int(nT+1)   #  round up to the nearest integer, the number of time steps
    dt = T/nT        #  adjust the time step down accordingly
    
    T = np.linspace( 0., T, nT+1)             # steps of size dt between 0 and T,
                                              # including 0 and T
    X = np.zeros( ( n, (nT+1)), np.float64)   #  allocate path array
    for i in range(n):
       X[i,0] = X0
    for k in range(nT):  # take a time step for all n paths at once
    
       X[:,k+1] = X[:,k] - a*X[:,k]*dt + rng.standard_normal(n)*sig*np.sqrt(dt)
       
    return (X,T)
       
        
#------------------ The main program ---------------------------------------------------

rng = default_rng()      # instantiate a bit generator

#    Parameters for the OU process being simulated

a   = .2     # mean reversion parameter
sig = 1.     # noise amplitude
X0  = 0.     # starting point for simulated paths
dt  = .001     # simulation time step size

PathPlotFile    = "PathPics.pdf"      # filename for a picture of some paths
DensityPlotFile = "OU_densities.pdf"  # filename for a PDF plots

#   Make a picture with a few paths

n_bold = 3                 # number of bold paths
n_soft = 50                # number of softer paths
T_f    = 2.                # simulate to this "fial" time

fig, ax = plt.subplots()     # Create a figure containing a single axes.

#   generate and plot the bold paths

X,T = path_sim( X0, n_bold, dt, T_f, a, sig, rng)

bold_colors = ['blue', 'green', 'red']
for i in range(n_bold):
   ax.plot( T, X[i,:], linewidth = 1, color = bold_colors[i])
   
#   generate and plot the soft paths

X,T = path_sim( X0, n_soft, dt, T_f, a, sig, rng)
for i in range(n_soft):
   ax.plot( T, X[i,:], linewidth = .1, color = 'gray')

plt.ylim( -2.5, 2.5)
plt.xlabel("T")
plt.ylabel("X")
title_string = r"Ornstein Uhlenbeck paths, $a=${a:6.2f}, $\sigma =${sig:6.2f}"
title_string = title_string.format(a=a, sig=sig)
plt.title(title_string)
plt.grid()
plt.savefig(PathPlotFile)                                    # save in the same directory
plt.show()


fig, ax = plt.subplots()     # Create a figure containing a single axes.

T1 = .5
T2 = 2.

n_hist = 100000
n_bins = 40
X_hist = []

X,T = path_sim( X0, n_hist, dt, T1, a, sig, rng)
hist1 = "T = {T1:6.2f}".format(T1=T1)
nn,nT = X.shape
X_hist.append( X[:,nT-1] )
X,T = path_sim( X0, n_hist, dt, T2, a, sig, rng)
hist2 = "T = {T2:6.2f}".format(T2=T2)

nn,nT = X.shape
X_hist.append( X[:,nT-1] )

ax.hist(X_hist, bins = n_bins, range = (-4., 4.), density = True, label = [hist1, hist2])
ax.legend()
ax.grid()
plt.xlabel("X")
plt.ylabel("density")
HistTitle = r'Histogram estimated PDF of $X_T$ using{n_hist:10.2e} paths'
HistTitle = HistTitle.format(n_hist=n_hist)
plt.title(HistTitle)

plt.savefig(DensityPlotFile)
plt.show()