import numpy as np
import matplotlib.pyplot as plt
from scipy.stats import norm

# Define parameters
K = 1.0
sigma = 0.2
T = 1
x = np.linspace(-0.5, 0.5, 100)  
tau = np.linspace(0.0001, T, 100)
X, TAU = np.meshgrid(x, tau)

def black_scholes(x, tau):
    d_plus = (1 / (sigma * np.sqrt(tau))) * (x + 0.5 * sigma**2 * tau + 0.5 * sigma**2 * tau)
    d_minus = (1 / (sigma * np.sqrt(tau))) * (x + 0.5 * sigma**2 * tau - 0.5 * sigma**2 * tau)
    u = K * (np.exp(x + 0.5 * sigma**2 * tau) * norm.cdf(d_plus) - norm.cdf(d_minus))
    return u

fig = plt.figure(figsize=(10, 10))
ax = fig.add_subplot(projection='3d')

K = 1.0
sigma = 0.05
num_walks = 10
starting_points = [-0.2, -0.1, 0, 0.1, 0.2]
num_starting_points = len(starting_points)
num_points = 1000
t_values = np.linspace(1e-4, 1, num_points)

brownian_walks = np.zeros((num_starting_points, num_walks, num_points))
for i in range(num_starting_points):
    for j in range(num_walks):
        dt = 1 / num_points
        dW = np.random.normal(0, np.sqrt(dt), num_points)
        brownian_walks[i, j] = np.cumsum(dW) * sigma + starting_points[i]

for i in range(num_starting_points):
    for j in range(num_walks):
        B_t = brownian_walks[i,j]
        S_t = np.exp(B_t)-1
        S_t = S_t[::-1]
        ax.plot(T-t_values, S_t, black_scholes(S_t, t_values), color='w', alpha=0.2)

ax.plot_surface(T-TAU, X, black_scholes(X, TAU), cmap='viridis')

ax.set_ylabel(r'$(S-C)}} )
ax.set_xlabel(r'$t}})
ax.set_zlabel(r'$V(S, t)}})
ax.view_init(elev=20, azim=0, roll=0)

plt.savefig("black_scholes_surface.svg")
plt.show()