# -*- coding: utf-8 -*-
"""
Stochastic differential equations :
        
- simulate Brownian motion
- compute stochstic integral int W dW
- solve by Euler-Maruyama scheme the Ornstein–Uhlenbeck SDE
 dX_t = \theta(\mu - X_t) dt + \sigma dW_t
- solve by Euler-Maruyama scheme the Log-normal SDE: 
    dY_t = \mu Y_t dt + \sigma Y_t dW_t

@author: Gabriel Turinici, april 2025
"""

import numpy as np
import matplotlib.pyplot as plt

# Parameters
T = 1.0         # Time horizon
N = 500         # Number of time steps
M = 1000        # Number of scenarios
dt = T / N
t = np.linspace(0, T, N+1)

# 1. Simulate M scenarios of Brownian motion W on [0, T]
# time dimension is first, the second dimension are samples
W = np.zeros((N+1, M))                                # initialize
dW=np.zeros((N,M))
#TODO HERE: construct brownian motion using increments dW
dW=None
W=None

# Plot a few Brownian motion paths
plt.figure(figsize=(10, 6))
for i in range(5):
    plt.plot(t, W[:, i], lw=1)
plt.title("Sample Paths of Brownian Motion")
plt.xlabel("Time")
plt.ylabel("W(t)")
plt.grid(True)
plt.show()

# TODO 
# 2. Compute the Itô integral $\int_0^T W dW_t$ using the Itô approximation
ito_integral = None # shape (M,)
print("First 5 values of the Itô integral \int W dW:")
print(ito_integral[:5])
plt.figure("integral",figsize=(4,4))
plt.hist(ito_integral,alpha=0.25,label="integral",bins=30)
plt.hist((W[-1,:]**2-T)/2.,alpha=0.25,label='analytic formula',bins=30)
plt.legend()
plt.show()

# 3. Ornstein–Uhlenbeck process: dX_t = \theta(\mu - X_t) dt + \sigma dW_t
theta = 2.0*5
mu = 0.5
sigma_ou = 0.3
X0 = 0.0

X_ou = np.zeros((N+1, M))                             # shape (time, scenarios)
X_ou[0] = X0

# TODO : implement numerical scheme
for i in range(N):
    X_ou[i+1] = None
    
plt.figure(figsize=(10, 6))
# Plot a few Ornstein–Uhlenbeck paths
for i in range(5):
    plt.plot(t, X_ou[:, i], lw=1)
plt.title("Sample Paths of Ornstein-Uhlenbeck Process")
plt.xlabel("Time")
plt.ylabel("X(t)")
plt.grid(True)
plt.show()

print("Final values of first 5 OU paths:")
print(X_ou[-1, :5])

# 4. Log-normal process: dY_t = \mu Y_t dt + \sigma Y_t dW_t
mu_ln = 0.1
sigma_ln = 0.2
Y0 = 1.0

Y_ln = np.zeros((N+1, M))                             # shape (time, scenarios)
Y_ln[0] = Y0

# TODO : implement numerical scheme
for i in range(N):
    Y_ln[i+1] = None

# Plot a few log-normal paths
plt.figure(figsize=(10, 6))
for i in range(5):
    plt.plot(t, Y_ln[:, i], lw=1)
plt.title("Sample Paths of Log-Normal Process")
plt.xlabel("Time")
plt.ylabel("Y(t)")
plt.grid(True)
plt.show()

print("Final values of first 5 log-normal paths:")
print(Y_ln[-1, :5])