"""
Fig 11: Two-Parameter Heatmap (Temperature × Signal Quality)
Shows interaction effects on TTE
"""

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

def compute_tte_surface(T_a_range, psi_range):
    """Compute TTE for grid of (T_a, psi) values"""
    
    T_grid, psi_grid = np.meshgrid(T_a_range, psi_range)
    tte_grid = np.zeros_like(T_grid)
    
    # Battery capacity
    Q_full = 2.78  # Ah
    V_avg = 3.8    # V
    E_total = Q_full * V_avg  # Wh
    
    # Baseline power components
    P_bg = 5.0
    P_scr = 8.0 * 0.3  # L = 0.3
    P_cpu = 35.0 * 0.4  # C = 0.4
    P_gps = 0.015  # Minimal
    
    for i in range(len(psi_range)):
        for j in range(len(T_a_range)):
            T_a = T_grid[i, j]
            psi = psi_grid[i, j]
            
            # Network power depends on signal quality (worse signal = more power)
            k_net_base = 7.0
            k_net = k_net_base * (1 + 2 * (1 - psi))
            P_net = k_net * 0.05  # N = 0.05
            
            # Temperature affects efficiency (cold reduces capacity, hot increases resistance)
            temp_factor = 1.0
            if T_a < 10:
                temp_factor = 0.8 + 0.02 * T_a  # Reduced capacity in cold
            elif T_a > 30:
                temp_factor = 1.0 + 0.01 * (T_a - 30)  # Increased losses in heat
            
            P_total = (P_bg + P_scr + P_cpu + P_net + P_gps) * temp_factor
            
            tte_grid[i, j] = E_total / P_total
    
    return T_grid, psi_grid, tte_grid

def make_figure(config):
    """Generate Fig 11: Two-Parameter Heatmap"""
    
    set_oprice_style()
    
    # Parameter ranges
    T_a_range = np.linspace(-10, 40, 50)  # °C
    psi_range = np.linspace(0.1, 1.0, 50)  # Signal quality (0=poor, 1=excellent)
    
    # Compute TTE surface
    T_grid, psi_grid, tte_grid = compute_tte_surface(T_a_range, psi_range)
    
    # Create figure
    fig, ax = plt.subplots(figsize=(12, 9))
    
    # Heatmap
    im = ax.contourf(T_grid, psi_grid, tte_grid, levels=20, cmap='RdYlGn', alpha=0.9)
    
    # Contour lines
    contours = ax.contour(T_grid, psi_grid, tte_grid, levels=10, colors='k', 
                          linewidths=0.5, alpha=0.3)
    ax.clabel(contours, inline=True, fontsize=8, fmt='%.1f h')
    
    # Colorbar
    cbar = fig.colorbar(im, ax=ax, label='Time to Empty (hours)')
    cbar.ax.tick_params(labelsize=9)
    
    # Mark baseline condition
    T_baseline = 25
    psi_baseline = 0.8
    ax.plot(T_baseline, psi_baseline, 'r*', markersize=20, 
            markeredgecolor='white', markeredgewidth=1.5, 
            label='Baseline Condition', zorder=5)
    
    # Mark danger zones
    # Cold + poor signal
    ax.add_patch(plt.Rectangle((-10, 0.1), 15, 0.3, 
                               fill=False, edgecolor='red', linewidth=2, 
                               linestyle='--', label='Critical Zone'))
    ax.text(-5, 0.25, 'Battery Killer\\n(Cold + Poor Signal)', 
           fontsize=9, color='darkred', ha='center',
           bbox=dict(boxstyle='round', facecolor='white', alpha=0.7))
    
    # Optimal zone
    ax.add_patch(plt.Rectangle((15, 0.7), 15, 0.25, 
                               fill=False, edgecolor='green', linewidth=2, 
                               linestyle='--', label='Optimal Zone'))
    ax.text(22.5, 0.825, 'Optimal\\n(Mild + Good Signal)', 
           fontsize=9, color='darkgreen', ha='center',
           bbox=dict(boxstyle='round', facecolor='white', alpha=0.7))
    
    # Labels and styling
    ax.set_xlabel('Ambient Temperature (°C)', fontsize=11)
    ax.set_ylabel('Signal Quality (Ψ)', fontsize=11)
    ax.set_title('TTE Sensitivity to Temperature and Signal Quality (Interaction Effect)', 
                 fontsize=12, fontweight='bold')
    ax.set_xlim(-10, 40)
    ax.set_ylim(0.1, 1.0)
    ax.legend(loc='upper left', framealpha=0.9, fontsize=9)
    
    # Add statistics
    tte_min = np.min(tte_grid)
    tte_max = np.max(tte_grid)
    tte_range = tte_max - tte_min
    
    stats_text = f'TTE Range:\\n' + \
                f'Min: {tte_min:.2f}h\\n' + \
                f'Max: {tte_max:.2f}h\\n' + \
                f'Spread: {tte_range:.2f}h ({tte_range/tte_max*100:.1f}%)'
    
    ax.text(0.98, 0.02, stats_text, transform=ax.transAxes,
            fontsize=9, verticalalignment='bottom', horizontalalignment='right',
            bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.8),
            family='monospace')
    
    # Save
    figure_dir = config.get('global', {}).get('figure_dir', 'figures')
    os.makedirs(figure_dir, exist_ok=True)
    output_base = os.path.join(figure_dir, 'Fig11_Heatmap_Temp_Signal')
    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_min_h": float(tte_min),
            "tte_max_h": float(tte_max),
            "tte_range_h": float(tte_range)
        },
        "validation_flags": {},
        "pass": True
    }