"""
Brownian motion implementations: 
- construct a brownian motion (BM), first
only a scenario then several realizations;
-then compute the stochastic integral $\int W d W$
 
    @author: Gabriel Turinici, april 2024
"""
#import libraries
import numpy as np
import matplotlib.pyplot as plt

print('construct a brownian motion realization called also "scenario"')
T=1.
N=500
h=T/N
#time range
trange= np.linspace(0,T,N+1,endpoint=True)

print('first attempt, incorrect correlation in W')
W= np.random.randn(N+1)*np.sqrt(np.linspace(0,T,num=N+1,endpoint=True))
plt.figure('incorrect correlation in W')
plt.plot(W)
#correlation struction is not good (is always zero !)


#construct with increments
dW = np.random.randn(N)*np.sqrt(h)#these are increments of B.M.
W=np.zeros(N+1)
W[1:]=np.cumsum(dW)
plt.figure('W')
plt.plot(W)

print('vectorized version with many realisations:', 
      'first index time, second=realizations')
M=5_000
dW = np.random.randn(N,M)*np.sqrt(h)
W=np.zeros((N+1,M))
W[1:,:]=np.cumsum(dW,axis=0)
plt.figure('W vect')
res=plt.plot(trange,W[:,:np.minimum(M,50)])#plot first 50 of them

print('start computing Riemann-Ito sum')
ri_sum = np.sum(W[:-1,:]*dW,axis=0) 

plt.figure('hist',figsize=(6,3))
exact = (W[-1,:]**2-T)/2#obtained by Ito formula
plt.subplot(1,2,1)
plt.hist(exact,bins=40,alpha=0.5,label='exact')
plt.hist(ri_sum,bins=40,alpha=0.5,label='Rieman-Ito sum')
plt.legend()
plt.subplot(1,2,2)
plt.hist(exact-ri_sum,bins=40,alpha=0.5,label='difference')
plt.legend()