import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation, PillowWriter

# --- 1. SETUP & SIMULATION DATA ---
np.random.seed(42)
n_frames = 150
x = np.arange(n_frames)

# Regime Transitions: 0-60 (Stable), 60-100 (Latent), 100-150 (Stress)
regime_bg = np.zeros(n_frames)
regime_bg[60:100] = 1  # Latent Build-up
regime_bg[100:] = 2   # Stress

# Simulated Signals
volatility = np.random.normal(0.5, 0.1, n_frames)
volatility[100:] += np.cumsum(np.random.normal(0.2, 0.05, 50)) # Reactive Spike

entropy = np.random.normal(0.2, 0.05, n_frames)
entropy[65:] += np.linspace(0, 0.6, 85) # Early build-up in latent phase

trigger_signal = (entropy > 0.45).astype(float) # The "Early" Detection

# --- 2. PLOTTING ENGINE ---
fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(10, 6), sharex=True)
plt.subplots_adjust(hspace=0.3)

# Style Customization
for ax in [ax1, ax2]:
    ax.set_facecolor('#0f0f0f')
    ax.grid(color='white', alpha=0.1)
    ax.tick_params(colors='white')

fig.patch.set_facecolor('#0f0f0f')

# Lines
line_vol, = ax1.plot([], [], color='#00ffcc', lw=2, label='Market Volatility (Reactive)')
line_ent, = ax2.plot([], [], color='#ff00ff', lw=2, label='HMM Entropy (Latent Signal)')
line_trig, = ax2.plot([], [], color='#ffcc00', lw=3, label='EARLY TRIGGER')

ax1.set_ylim(0, 3)
ax2.set_ylim(0, 1.2)
ax1.legend(loc='upper left')
ax2.legend(loc='upper left')

# --- 3. ANIMATION LOGIC ---
def init():
    line_vol.set_data([], [])
    line_ent.set_data([], [])
    line_trig.set_data([], [])
    return line_vol, line_ent, line_trig

def update(frame):
    # Update line data
    line_vol.set_data(x[:frame], volatility[:frame])
    line_ent.set_data(x[:frame], entropy[:frame])
    
    if frame > 75: # Trigger activation point
        line_trig.set_data(x[75:frame], trigger_signal[75:frame])
    
    # Dynamic Shading for Regimes
    if frame == 60:
        ax1.axvspan(60, 100, color='yellow', alpha=0.1, label='Latent Phase')
    if frame == 100:
        ax1.axvspan(100, 150, color='red', alpha=0.2, label='Stress Phase')
        
    return line_vol, line_ent, line_trig

# --- 4. EXPORT ---
ani = FuncAnimation(fig, update, frames=n_frames, init_func=init, blit=True)

print("Encoding GIF... this may take a moment.")
writer = PillowWriter(fps=20)
ani.save("assets/detection.gif", writer=writer)
print("GIF saved to assets/detection.gif")