"""
Fig 13: Survival/Reliability Curve
Shows probability of battery lasting beyond time t
"""

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

def make_figure(config):
    """Generate Fig 13: Survival Curve"""
    
    set_oprice_style()
    
    seed = config.get('global', {}).get('seed', 42)
    np.random.seed(seed)
    
    # Generate TTE distribution (log-normal)
    n_samples = 300
    tte_mean = 2.5  # hours
    tte_std = 0.4
    
    # Log-normal parameters
    mu = np.log(tte_mean**2 / np.sqrt(tte_mean**2 + tte_std**2))
    sigma = np.sqrt(np.log(1 + tte_std**2 / tte_mean**2))
    
    tte_samples = np.random.lognormal(mu, sigma, n_samples)
    tte_samples = np.clip(tte_samples, 1.0, 5.0)
    
    # Sort for survival function
    tte_sorted = np.sort(tte_samples)
    
    # Compute survival function: S(t) = P(TTE > t)
    t_values = np.linspace(0, 4.5, 200)
    survival_prob = np.zeros(len(t_values))
    
    for i, t in enumerate(t_values):
        survival_prob[i] = np.sum(tte_sorted > t) / n_samples
    
    # Find confidence levels
    tte_50 = np.percentile(tte_sorted, 50)  # Median
    tte_95 = np.percentile(tte_sorted, 5)   # 95% confidence (5th percentile)
    
    # Create figure
    fig, ax = plt.subplots(figsize=(12, 8))
    
    # Plot survival curve
    ax.plot(t_values, survival_prob * 100, 'b-', linewidth=3, label='Survival Function')
    
    # Fill area under curve
    ax.fill_between(t_values, 0, survival_prob * 100, alpha=0.2, color='lightblue')
    
    # Mark key percentiles
    ax.axhline(50, color='orange', linestyle='--', linewidth=1.5, alpha=0.7, label='50% Survival')
    ax.axvline(tte_50, color='orange', linestyle=':', linewidth=1.5, alpha=0.7)
    ax.plot(tte_50, 50, 'o', markersize=10, color='orange', markeredgecolor='white', 
            markeredgewidth=2, zorder=5)
    ax.text(tte_50 + 0.1, 50, f'Median TTE = {tte_50:.2f}h', fontsize=10, 
            verticalalignment='center')
    
    ax.axhline(95, color='green', linestyle='--', linewidth=1.5, alpha=0.7, label='95% Confidence')
    ax.axvline(tte_95, color='green', linestyle=':', linewidth=1.5, alpha=0.7)
    ax.plot(tte_95, 95, 's', markersize=10, color='green', markeredgecolor='white', 
            markeredgewidth=2, zorder=5)
    ax.text(tte_95 + 0.1, 95, f'95% Confident TTE = {tte_95:.2f}h', fontsize=10, 
            verticalalignment='center')
    
    # Annotate interpretation
    ax.annotate('95% of devices will\nlast at least this long',
                xy=(tte_95, 95), xytext=(tte_95 - 0.5, 75),
                arrowprops=dict(arrowstyle='->', color='green', lw=2),
                fontsize=10, color='darkgreen',
                bbox=dict(boxstyle='round', facecolor='lightgreen', alpha=0.7))
    
    # Labels and styling
    ax.set_xlabel('Time (hours)', fontsize=11)
    ax.set_ylabel('Survival Probability (%)', fontsize=11)
    ax.set_title('Battery Reliability: Survival Function Analysis', 
                 fontsize=12, fontweight='bold')
    ax.set_xlim(0, 4.5)
    ax.set_ylim(0, 105)
    ax.grid(True, alpha=0.3)
    ax.legend(loc='upper right', framealpha=0.9, fontsize=10)
    
    # Add distribution statistics
    stats_text = f'TTE Distribution:\\n'
    stats_text += f'Mean: {tte_mean:.2f}h\\n'
    stats_text += f'Std: {tte_std:.2f}h\\n'
    stats_text += f'Median: {tte_50:.2f}h\\n'
    stats_text += f'95% CI: {tte_95:.2f}h\\n'
    stats_text += f'Sample Size: {n_samples}'
    
    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 interpretation guide
    guide_text = 'Interpretation:\\n' + \
                 '• S(t) = P(TTE > t)\\n' + \
                 '• Steep drop = high variability\\n' + \
                 '• Use 95% value for conservative estimates'
    
    ax.text(0.98, 0.02, guide_text, transform=ax.transAxes,
            fontsize=9, verticalalignment='bottom', 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, 'Fig13_Survival_Curve')
    save_figure(fig, output_base, dpi=config.get('global', {}).get('dpi', 300))
    plt.close()
    
    # Validation
    is_monotonic = check_monotonic_decreasing(survival_prob)
    
    return {
        "output_files": [f"{output_base}.pdf", f"{output_base}.png"],
        "computed_metrics": {
            "median_tte_h": float(tte_50),
            "tte_95_confidence_h": float(tte_95),
            "n_samples": n_samples
        },
        "validation_flags": {"survival_monotonic": is_monotonic},
        "pass": is_monotonic
    }