"""
Fig 4: Internal Resistance 3D Surface
Shows R0(T, z) dependency on temperature and SOC
"""

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

def compute_internal_resistance(T_b, z, R_ref, E_a, T_ref):
    """
    Compute internal resistance with temperature and SOC dependence
    R0 = R_ref * exp(E_a/R * (1/T - 1/T_ref)) * (1 + k_z * (1-z))
    """
    R_gas = 8.314  # J/(mol·K)
    k_z = 0.5      # SOC dependence coefficient
    
    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 make_figure(config):
    """Generate Fig 4: Internal Resistance 3D Surface"""
    
    set_oprice_style()
    
    # Get parameters
    params = config.get('battery_params', {})
    R_ref = params.get('R_ref', 0.1)
    E_a = params.get('E_a', 20000)
    T_ref = params.get('T_ref', 298.15)
    
    # Create grid
    T_celsius = np.linspace(-10, 40, 50)
    T_kelvin = T_celsius + 273.15
    z_values = np.linspace(0.05, 0.95, 50)
    T_grid, z_grid = np.meshgrid(T_kelvin, z_values)
    
    # Compute resistance surface
    R0_grid = compute_internal_resistance(T_grid, z_grid, R_ref, E_a, T_ref)
    
    # Create figure
    fig = plt.figure(figsize=(12, 9))
    ax = fig.add_subplot(111, projection='3d')
    
    # Plot surface
    surf = ax.plot_surface(T_celsius, z_grid, R0_grid, 
                           cmap='coolwarm', alpha=0.9, 
                           edgecolor='none', antialiased=True)
    
    # Add contour lines on bottom
    ax.contour(T_celsius, z_values, R0_grid, 
               levels=10, offset=0, cmap='coolwarm', alpha=0.5)
    
    # Labels and title
    ax.set_xlabel('Temperature (°C)', fontsize=11, labelpad=10)
    ax.set_ylabel('State of Charge (SOC)', fontsize=11, labelpad=10)
    ax.set_zlabel('Internal Resistance (Ω)', fontsize=11, labelpad=10)
    ax.set_title('Internal Resistance Dependence on Temperature and SOC', 
                 fontsize=12, fontweight='bold', pad=20)
    
    # Set viewing angle
    ax.view_init(elev=20, azim=135)
    
    # Colorbar
    cbar = fig.colorbar(surf, ax=ax, shrink=0.6, aspect=15, pad=0.1)
    cbar.set_label('R₀ (Ω)', fontsize=10)
    
    # Add annotation for key regions
    ax.text2D(0.02, 0.95, 'Key Observations:\n' + 
              '• Low temp → High resistance\n' +
              '• Low SOC → High resistance\n' +
              '• Coupled effect at extremes',
              transform=ax.transAxes, fontsize=9,
              verticalalignment='top',
              bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.8))
    
    # Save
    figure_dir = config.get('global', {}).get('figure_dir', 'figures')
    os.makedirs(figure_dir, exist_ok=True)
    output_base = os.path.join(figure_dir, 'Fig04_Internal_Resistance')
    save_figure(fig, output_base, dpi=config.get('global', {}).get('dpi', 300))
    plt.close()
    
    # Compute some statistics
    R0_min = float(np.min(R0_grid))
    R0_max = float(np.max(R0_grid))
    R0_ratio = R0_max / R0_min
    
    return {
        "output_files": [f"{output_base}.pdf", f"{output_base}.png"],
        "computed_metrics": {
            "R0_min_ohm": R0_min,
            "R0_max_ohm": R0_max,
            "ratio": R0_ratio
        },
        "validation_flags": {},
        "pass": True
    }