"""
Fig 12: Monte Carlo Spaghetti Plot
Shows uncertainty in TTE predictions with random perturbations
"""

import os
import numpy as np
import matplotlib.pyplot as plt
from plot_style import set_oprice_style, save_figure

def simulate_stochastic_trajectory(base_power, duration_hours, n_points, 
                                   sigma_P, theta_P, dt, seed_offset=0):
    """
    Simulate one stochastic SOC trajectory with OU process for power
    
    Args:
        base_power: Base power consumption (W)
        duration_hours: Simulation duration
        n_points: Number of time points
        sigma_P: OU process volatility
        theta_P: OU process mean reversion rate
        dt: Time step (hours)
        seed_offset: Random seed offset
    """
    
    np.random.seed(42 + seed_offset)
    
    # Time axis
    t_h = np.linspace(0, duration_hours, n_points)
    
    # Ornstein-Uhlenbeck process for power fluctuation
    dW = np.random.normal(0, np.sqrt(dt), n_points)
    P_perturbation = np.zeros(n_points)
    
    for i in range(1, n_points):
        P_perturbation[i] = P_perturbation[i-1] * (1 - theta_P * dt) + sigma_P * dW[i]
    
    P_total = base_power + P_perturbation
    P_total = np.clip(P_total, base_power * 0.5, base_power * 1.5)
    
    # Compute SOC trajectory
    Q_full = 2.78  # Ah
    V_avg = 3.8    # V
    E_total = Q_full * V_avg  # Wh
    
    z = np.zeros(n_points)
    z[0] = 1.0
    
    for i in range(1, n_points):
        # Energy consumed in this step
        dE = P_total[i] * dt
        dz = -dE / E_total
        z[i] = z[i-1] + dz
        
        if z[i] < 0.05:
            z[i] = 0.05
            break
    
    # Fill remaining with cutoff value
    z[z < 0.05] = 0.05
    
    return t_h, z

def make_figure(config):
    """Generate Fig 12: Monte Carlo Spaghetti Plot"""
    
    set_oprice_style()
    
    # Simulation parameters
    n_paths = 100
    duration_hours = 4.0
    n_points = 200
    dt = duration_hours / n_points
    
    # Baseline power
    base_power = 15.0  # W
    
    # OU process parameters for power uncertainty
    sigma_P = 2.0   # Volatility
    theta_P = 0.5   # Mean reversion rate
    
    # Create figure
    fig, ax = plt.subplots(figsize=(12, 8))
    
    # Store all paths for statistics
    all_z = []
    all_tte = []
    
    # Simulate and plot paths
    for i in range(n_paths):
        t_h, z = simulate_stochastic_trajectory(base_power, duration_hours, n_points,
                                                sigma_P, theta_P, dt, seed_offset=i)
        
        # Plot path
        if i < 50:  # Only plot half to avoid overcrowding
            ax.plot(t_h, z * 100, color='gray', alpha=0.15, linewidth=0.8, zorder=1)
        
        all_z.append(z)
        
        # Find TTE (first time z <= 0.05)
        cutoff_idx = np.where(z <= 0.05)[0]
        if len(cutoff_idx) > 0:
            tte = t_h[cutoff_idx[0]]
        else:
            tte = duration_hours
        all_tte.append(tte)
    
    # Compute statistics
    all_z = np.array(all_z)
    z_mean = np.mean(all_z, axis=0)
    z_std = np.std(all_z, axis=0)
    z_p5 = np.percentile(all_z, 5, axis=0)
    z_p95 = np.percentile(all_z, 95, axis=0)
    
    t_h_ref = np.linspace(0, duration_hours, n_points)
    
    # Plot confidence band
    ax.fill_between(t_h_ref, z_p5 * 100, z_p95 * 100, 
                    alpha=0.3, color='lightblue', label='90% Confidence Band', zorder=2)
    
    # Plot mean
    ax.plot(t_h_ref, z_mean * 100, 'b-', linewidth=3, label='Mean Trajectory', zorder=3)
    
    # Plot ±1 std
    ax.plot(t_h_ref, (z_mean + z_std) * 100, 'b--', linewidth=1.5, alpha=0.6, 
            label='Mean ± 1σ', zorder=2)
    ax.plot(t_h_ref, (z_mean - z_std) * 100, 'b--', linewidth=1.5, alpha=0.6, zorder=2)
    
    # Cutoff line
    ax.axhline(5, color='r', linestyle='--', linewidth=1.5, alpha=0.7, label='Cutoff (5%)')
    
    # Labels and styling
    ax.set_xlabel('Time (hours)', fontsize=11)
    ax.set_ylabel('State of Charge (%)', fontsize=11)
    ax.set_title('Monte Carlo Uncertainty Quantification (N=100 paths)', 
                 fontsize=12, fontweight='bold')
    ax.set_xlim(0, duration_hours)
    ax.set_ylim(0, 105)
    ax.grid(True, alpha=0.3)
    ax.legend(loc='upper right', framealpha=0.9, fontsize=10)
    
    # Add statistics box
    tte_mean = np.mean(all_tte)
    tte_std = np.std(all_tte)
    tte_p5 = np.percentile(all_tte, 5)
    tte_p95 = np.percentile(all_tte, 95)
    
    stats_text = f'TTE Statistics (hours):\\n'
    stats_text += f'Mean: {tte_mean:.2f} ± {tte_std:.2f}\\n'
    stats_text += f'5th %ile: {tte_p5:.2f}\\n'
    stats_text += f'95th %ile: {tte_p95:.2f}\\n'
    stats_text += f'Range: [{tte_p5:.2f}, {tte_p95:.2f}]\\n'
    stats_text += f'Coefficient of Variation: {tte_std/tte_mean*100:.1f}%'
    
    ax.text(0.02, 0.35, stats_text, transform=ax.transAxes,
            fontsize=9, verticalalignment='top',
            bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.8),
            family='monospace')
    
    # Add annotation explaining OU process
    ax.text(0.98, 0.97, 'Uncertainty Model:\\nOrnstein-Uhlenbeck\\nprocess for power\\nfluctuations',
            transform=ax.transAxes, fontsize=9,
            verticalalignment='top', horizontalalignment='right',
            bbox=dict(boxstyle='round', facecolor='lightblue', alpha=0.7))
    
    plt.tight_layout()
    
    # Save
    figure_dir = config.get('global', {}).get('figure_dir', 'figures')
    os.makedirs(figure_dir, exist_ok=True)
    output_base = os.path.join(figure_dir, 'Fig12_Monte_Carlo')
    save_figure(fig, output_base, dpi=config.get('global', {}).get('dpi', 300))
    plt.close()
    
    return {
        "output_files": [f"{output_base}.pdf", f"{output_base}.png"],
        "computed_metrics": {
            "tte_mean_h": float(tte_mean),
            "tte_std_h": float(tte_std),
            "tte_cv_pct": float(tte_std/tte_mean * 100),
            "n_paths": n_paths
        },
        "validation_flags": {},
        "pass": True
    }