"""
Figure 9 from the thesis:
Nested simulation - outer scenarios [0, T0] followed by inner scenarios [T0, T1].
"""

import numpy as np
import matplotlib.pyplot as plt
from core import sim_brownian_motion, exact_gbm

np.random.seed(42)

s0 = 100.0
mu = 0.05
sigma = 0.2
T0 = 1.0
T1 = 2.0
l = 10
N = 3   # outer scenarios
M = 10  # inner scenarios per outer

fig, ax = plt.subplots(figsize=(10, 6))

# Outer simulation: GBM paths from t=0 to T0
bm_outer = sim_brownian_motion(T0, l, N)
t_outer = np.linspace(0, T0, 2**l + 1)

gbm_outer = np.zeros((2**l + 1, N))
for n in range(N):
    gbm_outer[:, n] = exact_gbm(s0, mu, sigma, t_outer, bm_outer[:, n])
    ax.plot(t_outer, gbm_outer[:, n], linewidth=1.0)

# Inner simulation: from each outer endpoint, simulate M inner paths
t_inner = np.linspace(T0, T1, 2**l + 1)

for n in range(N):
    s0_inner = gbm_outer[-1, n]
    bm_inner = sim_brownian_motion(T1 - T0, l, M)
    for m in range(M):
        gbm_inner = exact_gbm(s0_inner, mu, sigma, t_inner - T0, bm_inner[:, m])
        ax.plot(t_inner, gbm_inner, linewidth=0.5, alpha=0.7)

# Vertical line at T0
ax.axvline(x=T0, color='cyan', linewidth=1.5, linestyle='-')

ax.set_xlabel('t')
ax.set_title('Nested simulation')
fig.tight_layout()
fig.savefig('plot_nested_simulation.png', dpi=150)
print("Saved: plot_nested_simulation.png")

plt.show()