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

n_points = 1000
t = np.linspace(0, 1, n_points)

def brownian_bridge(n):
    dw = np.random.normal(0, np.sqrt(1/n), n)
    w = np.cumsum(dw)
    bb = w - t * w[-1]
    return bb

bridge = brownian_bridge(n_points)

tau_star_idx = np.argmin(bridge)
tau_star = t[tau_star_idx]
min_value = bridge[tau_star_idx]

# Create the Brownian excursion by switching paths and shifting
first_part = bridge[:tau_star_idx+1]
second_part = bridge[tau_star_idx:]

# Transform to excursion
excursion = np.concatenate([
    second_part - min_value,
    first_part[1:] - min_value
])

# Create custom colormap for visualization
colors_bridge = np.ones((n_points, 4))  # RGBA
colors_bridge[:tau_star_idx+1, 0:2] = 0, 1  # Blue for [0, tau_star]
colors_bridge[tau_star_idx:, 1:3] = 1, 0    # Green for [tau_star, 1]

colors_excursion = np.ones((n_points, 4))  # RGBA
colors_excursion[:n_points-tau_star_idx, 1:3] = 1, 0    # Green first
colors_excursion[n_points-tau_star_idx:, 0:2] = 0, 1    # Blue second

# Create the figure
plt.figure(figsize=(10, 8))
gs = gridspec.GridSpec(2, 1, height_ratios=[1, 1])

# Plot Brownian bridge
ax1 = plt.subplot(gs[0])
ax1.plot(t, bridge, '-', linewidth=1.5)
ax1.plot(t[:tau_star_idx+1], bridge[:tau_star_idx+1], 'b-', linewidth=1.5)
ax1.plot(t[tau_star_idx:], bridge[tau_star_idx:], 'g-', linewidth=1.5)
ax1.axhline(y=0, color='black', linestyle='-', alpha=0.3)
ax1.axvline(x=t[tau_star_idx], color='black', linestyle='--', alpha=0.3)
ax1.axhline(y=min_value, color='black', linestyle='--', alpha=0.3)
ax1.set_title('Brownian Bridge', fontsize=14)
ax1.set_xlim(0, 1)
# ax1.grid(False)

# Plot Brownian excursion
ax2 = plt.subplot(gs[1])
ax2.plot(t, excursion, '-', linewidth=1.5)
green_length = len(second_part)
ax2.plot(t[:green_length], excursion[:green_length], 'g-', linewidth=1.5)
ax2.plot(t[green_length-1:], excursion[green_length-1:], 'b-', linewidth=1.5)
ax2.axvline(x=t[green_length], color='black', linestyle='--', alpha=0.3)
ax2.set_title('Brownian Excursion', fontsize=14)
ax2.set_xlim(0, 1)
# ax2.grid(False)

plt.tight_layout()
plt.savefig('brownian_bridge_to_excursion.svg', bbox_inches='tight')
plt.show()