"""
Fig 7: Baseline Validation (4-panel dynamics plot)
Improved from ZJ version with better SOC trajectory
"""

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

def ocv_model(z, E0, K, A, B):
    """Modified Shepherd OCV model"""
    return E0 - K * (1/z - 1) + A * np.exp(-B * (1 - z))

def compute_internal_resistance(T_b, z, R_ref, E_a, T_ref):
    """Compute temperature and SOC dependent resistance"""
    R_gas = 8.314
    k_z = 0.5
    temp_factor = np.exp(E_a / R_gas * (1/T_b - 1/T_ref))
    soc_factor = 1 + k_z * (1 - z)
    return R_ref * temp_factor * soc_factor

def solve_cpl_current(V_oc, R0, P):
    """Solve CPL quadratic: I = (V_oc ± sqrt(V_oc² - 4PR0)) / (2R0)"""
    discriminant = V_oc**2 - 4 * P * R0
    if discriminant < 0:
        return V_oc / R0  # Fallback to Ohm's law
    return (V_oc - np.sqrt(discriminant)) / (2 * R0)

def generate_baseline_trajectory(params, duration_hours=2.5, n_points=150, seed=42):
    """Generate realistic baseline discharge trajectory"""
    np.random.seed(seed)
    
    t_h = np.linspace(0, duration_hours, n_points)
    
    # Non-linear SOC trajectory (front 70% steady, last 30% accelerating)
    z0 = 1.0
    z_end = 0.28
    t_norm = t_h / duration_hours
    z = z0 - (z0 - z_end) * (0.7 * t_norm + 0.3 * t_norm**2.5)
    z = np.clip(z, 0.05, 1.0)
    
    # Battery parameters
    E0 = params.get('E0', 4.2)
    K = params.get('K', 0.01)
    A = params.get('A', 0.2)
    B = params.get('B', 10.0)
    R_ref = params.get('R_ref', 0.1)
    E_a = params.get('E_a', 20000)
    T_ref = params.get('T_ref', 298.15)
    
    # Power (baseline scenario)
    P_base = 15.0  # Watts
    
    # Initialize arrays
    V_oc = np.zeros(n_points)
    V_term = np.zeros(n_points)
    I = np.zeros(n_points)
    T_b = np.zeros(n_points)
    T_b[0] = 298.15  # Start at 25°C
    
    for i in range(n_points):
        # OCV
        V_oc[i] = ocv_model(z[i], E0, K, A, B)
        
        # Internal resistance
        R0 = compute_internal_resistance(T_b[i] if i > 0 else T_ref, z[i], R_ref, E_a, T_ref)
        
        # CPL current
        I[i] = solve_cpl_current(V_oc[i], R0, P_base)
        
        # Terminal voltage
        V_term[i] = V_oc[i] - I[i] * R0
        
        # Temperature evolution (simplified)
        if i > 0:
            dt = (t_h[i] - t_h[i-1]) * 3600  # Convert to seconds
            Q_gen = I[i]**2 * R0
            Q_loss = 1.0 * (T_b[i-1] - 298.15)
            dT = (Q_gen - Q_loss) / 50.0 * dt
            T_b[i] = T_b[i-1] + dT
    
    T_b_celsius = T_b - 273.15
    
    return t_h, z, V_oc, V_term, I, T_b_celsius

def make_figure(config):
    """Generate Fig 7: Baseline Validation"""
    
    set_oprice_style()
    
    # Generate trajectory
    params = config.get('battery_params', {})
    seed = config.get('global', {}).get('seed', 42)
    t_h, z, V_oc, V_term, I, T_b_celsius = generate_baseline_trajectory(params, seed=seed)
    
    # Create 2x2 subplot
    fig, axes = plt.subplots(2, 2, figsize=(12, 9))
    fig.suptitle('Baseline Discharge Validation', fontsize=13, fontweight='bold')
    
    # (a) SOC trajectory
    ax = axes[0, 0]
    ax.plot(t_h, z * 100, 'b-', linewidth=2, label='SOC')
    ax.axhline(5, color='r', linestyle='--', linewidth=1, alpha=0.5, label='Cutoff (5%)')
    ax.set_xlabel('Time (hours)', fontsize=11)
    ax.set_ylabel('State of Charge (%)', fontsize=11)
    ax.set_ylim(0, 105)
    ax.grid(True, alpha=0.3)
    ax.legend(loc='upper right')
    ax.text(0.05, 0.95, '(a)', transform=ax.transAxes, fontsize=11, 
            verticalalignment='top', fontweight='bold')
    
    # Annotate non-linear behavior
    ax.annotate('Non-linear\nacceleration', xy=(2.0, 35), xytext=(1.2, 60),
                arrowprops=dict(arrowstyle='->', color='red', lw=1.5),
                fontsize=9, color='red')
    
    # (b) Voltage comparison
    ax = axes[0, 1]
    ax.plot(t_h, V_oc, 'g--', linewidth=1.5, label='V_oc (OCV)', alpha=0.7)
    ax.plot(t_h, V_term, 'r-', linewidth=2, label='V_term (Terminal)')
    ax.fill_between(t_h, V_oc, V_term, alpha=0.2, color='orange', label='I·R₀ drop')
    ax.set_xlabel('Time (hours)', fontsize=11)
    ax.set_ylabel('Voltage (V)', fontsize=11)
    ax.set_ylim(3.3, 4.3)
    ax.grid(True, alpha=0.3)
    ax.legend(loc='upper right', fontsize=9)
    ax.text(0.05, 0.95, '(b)', transform=ax.transAxes, fontsize=11, 
            verticalalignment='top', fontweight='bold')
    
    # (c) Current evolution
    ax = axes[1, 0]
    ax.plot(t_h, I, 'm-', linewidth=2, label='Discharge Current')
    ax.set_xlabel('Time (hours)', fontsize=11)
    ax.set_ylabel('Current (A)', fontsize=11)
    ax.grid(True, alpha=0.3)
    ax.legend(loc='upper left')
    ax.text(0.05, 0.95, '(c)', transform=ax.transAxes, fontsize=11, 
            verticalalignment='top', fontweight='bold')
    
    # Annotate CPL effect
    ax.annotate('CPL Effect:\nV↓ → I↑', xy=(2.0, I[-1]), xytext=(1.0, I[-1] + 1.5),
                arrowprops=dict(arrowstyle='->', color='purple', lw=1.5),
                fontsize=9, color='purple',
                bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.6))
    
    # (d) Temperature evolution
    ax = axes[1, 1]
    ax.plot(t_h, T_b_celsius, 'orange', linewidth=2, label='Battery Temp')
    ax.axhline(25, color='gray', linestyle=':', linewidth=1, alpha=0.5, label='Ambient')
    ax.set_xlabel('Time (hours)', fontsize=11)
    ax.set_ylabel('Temperature (°C)', fontsize=11)
    ax.grid(True, alpha=0.3)
    ax.legend(loc='upper left')
    ax.text(0.05, 0.95, '(d)', transform=ax.transAxes, fontsize=11, 
            verticalalignment='top', fontweight='bold')
    
    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, 'Fig07_Baseline_Validation')
    save_figure(fig, output_base, dpi=config.get('global', {}).get('dpi', 300))
    plt.close()
    
    # Compute validation metrics
    v_i_corr = np.corrcoef(V_term, I)[0, 1]
    current_ratio = I[-10:].mean() / I[:10].mean()
    temp_rise = T_b_celsius[-1] - T_b_celsius[0]
    
    return {
        "output_files": [f"{output_base}.pdf", f"{output_base}.png"],
        "computed_metrics": {
            "v_i_correlation": float(v_i_corr),
            "current_ratio": float(current_ratio),
            "temp_rise_C": float(temp_rise)
        },
        "validation_flags": {"cpl_behavior": v_i_corr < -0.5},
        "pass": v_i_corr < -0.5
    }