import numpy as np
import scipy.stats as st 
import matplotlib.pyplot as plt 

np.random.seed(42)

#defines the standard random walk
def RW(steps = 100, a = 1./2):
    assert a <= 1./2, 'Parameter a should be at most 0.5.'
    N = steps
    X = np.zeros(N+1)
    for i in range(1, N+1):
        U = st.uniform.rvs()
        X[i] = X[i-1] + (U <= a) - (a < U) * (U <= 2*a) 
    return X

#defines the rescaled random walk
def rescaledRW(T = 1., n = 100, a = 0.2):
    X = RW(n, a)
    dt = T / (n+1)
    return np.sqrt(dt / (2*a)) * X

#defines the Wiener process
def Wiener(T = 1., n = 100):
    W = np.zeros(n+1)
    dt = T / (n+1)
    for i in range(1,n+1):
        W[i] = W[i-1] + np.sqrt(dt) * st.norm.rvs()
    return W

#X = rescaledRW(n = 1000, a = 0.2)
#W = Wiener(n = 1000)
fig, axes = plt.subplots(nrows = 1, ncols = 3, figsize = (12,4))
i = 1
for ax in axes:
    n = 100*(10**i)
    print('running for n = ', n)
    X = rescaledRW(n = n, a = 0.2)
    W = Wiener(n = n)
    t = np.linspace(0., 1, n+1)
    ax.plot(t, X, linewidth = 1, label = 'Rescaled RW')
    ax.plot(t, W, linewidth = 1, label = 'Wiener')
    ax.set_title('time steps = '+str(n))
    i = i + 1
plt.legend()
plt.show()