import numpy as np
import matplotlib.pyplot as plt
from matplotlib.colors import LinearSegmentedColormap
from matplotlib.collections import LineCollection
import matplotlib.patches as mpatches

n_points = 5000
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)

zero_crossings = np.where(np.diff(np.signbit(bridge)))[0]
segments = np.concatenate([[0], zero_crossings, [n_points-1]])
segments_pairs = [(segments[i], segments[i+1]) for i in range(len(segments)-1)]

plt.figure(figsize=(12, 7))

points = np.array([t, bridge]).T.reshape(-1, 1, 2)
segments_lc = np.concatenate([points[:-1], points[1:]], axis=1)

colors = np.zeros(len(segments_lc))

for i, (start, end) in enumerate(segments_pairs):
    is_positive = bridge[start:end+1].mean() > 0
    colors[start:end] = 1 if is_positive else 0
    
    x_segment = t[start:end+1]
    y_segment = bridge[start:end+1]
    if is_positive:
        plt.fill_between(x_segment, y_segment, 0, alpha=0.3, color='green', label='Positive excursion' if i == 0 else "")
    else:
        plt.fill_between(x_segment, y_segment, 0, alpha=0.3, color='red', label='Negative excursion' if i == 1 else "")

cmap = LinearSegmentedColormap.from_list('pos_neg', ['red', 'green'])
lc = LineCollection(segments_lc, cmap=cmap, array=colors, linewidth=1, alpha=0.4)
plt.gca().add_collection(lc)

plt.axhline(y=0, color='black', linestyle='-', alpha=0.7)

green_patch = mpatches.Patch(color='green', alpha=0.3, label='Positive excursions')
red_patch = mpatches.Patch(color='red', alpha=0.3, label='Negative excursions')
plt.title('Brownian bridge as a union of excursions', fontsize=16)
plt.xlim(0, 1)
plt.grid(False)

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