"""
Fig 5: Radio Tail Energy Illustration
Shows the tail effect in network power consumption
"""

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

def make_figure(config):
    """Generate Fig 5: Radio Tail Energy Illustration"""
    
    set_oprice_style()
    
    # Time axis
    t = np.linspace(0, 10, 1000)
    
    # Data burst events (brief pulses)
    burst_times = [1.0, 4.5, 7.0]
    burst_duration = 0.2
    data_activity = np.zeros_like(t)
    
    for bt in burst_times:
        mask = (t >= bt) & (t < bt + burst_duration)
        data_activity[mask] = 1.0
    
    # Power state with tail effect
    # After each burst, power decays exponentially
    power_state = np.zeros_like(t)
    tau_tail = 2.0  # Tail decay time constant
    
    for i, ti in enumerate(t):
        # Find most recent burst
        recent_bursts = [bt for bt in burst_times if bt <= ti]
        if recent_bursts:
            t_since_burst = ti - max(recent_bursts)
            if t_since_burst < burst_duration:
                # During burst: high power
                power_state[i] = 1.0
            else:
                # After burst: exponential decay (tail)
                power_state[i] = 1.0 * np.exp(-(t_since_burst - burst_duration) / tau_tail)
    
    # Create figure with two subplots
    fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(10, 7), sharex=True)
    
    # Top: Data activity
    ax1.fill_between(t, 0, data_activity, alpha=0.6, color='#2ca02c', label='Data Transmission')
    ax1.set_ylabel('Data Activity', fontsize=11)
    ax1.set_ylim(-0.1, 1.2)
    ax1.set_yticks([0, 1])
    ax1.set_yticklabels(['Idle', 'Active'])
    ax1.grid(True, alpha=0.3, axis='x')
    ax1.legend(loc='upper right')
    ax1.set_title('Network Radio Tail Effect Illustration', fontsize=12, fontweight='bold')
    
    # Annotate burst durations
    for bt in burst_times:
        ax1.annotate('', xy=(bt + burst_duration, 1.1), xytext=(bt, 1.1),
                    arrowprops=dict(arrowstyle='<->', color='black', lw=1))
        ax1.text(bt + burst_duration/2, 1.15, f'{int(burst_duration*1000)}ms',
                ha='center', fontsize=8)
    
    # Bottom: Power state
    ax2.fill_between(t, 0, power_state, alpha=0.6, color='#ff7f0e', label='Radio Power State')
    ax2.plot(t, power_state, 'r-', linewidth=1.5)
    ax2.set_xlabel('Time (seconds)', fontsize=11)
    ax2.set_ylabel('Power State', fontsize=11)
    ax2.set_ylim(-0.1, 1.2)
    ax2.set_yticks([0, 0.5, 1])
    ax2.set_yticklabels(['Idle', 'Mid', 'High'])
    ax2.grid(True, alpha=0.3)
    ax2.legend(loc='upper right')
    
    # Highlight tail regions
    for bt in burst_times:
        tail_start = bt + burst_duration
        tail_end = tail_start + 3 * tau_tail
        ax2.axvspan(tail_start, min(tail_end, 10), alpha=0.2, color='yellow')
    
    # Add annotation explaining the tail
    ax2.text(0.98, 0.95, 
             'Tail Effect:\nPower remains elevated\nafter data transmission\n' +
             r'$P(t) = P_{high} \cdot e^{-t/\tau}$',
             transform=ax2.transAxes,
             fontsize=9, verticalalignment='top', horizontalalignment='right',
             bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.8))
    
    # Add decay time constant annotation
    bt = burst_times[0]
    t_tau = bt + burst_duration + tau_tail
    idx_tau = np.argmin(np.abs(t - t_tau))
    ax2.plot([bt + burst_duration, t_tau], [1.0, power_state[idx_tau]], 
             'k--', linewidth=1, alpha=0.5)
    ax2.text(t_tau, power_state[idx_tau] + 0.05, r'$\tau$ = 2s', 
             fontsize=9, ha='left')
    
    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, 'Fig05_Radio_Tail')
    save_figure(fig, output_base, dpi=config.get('global', {}).get('dpi', 300))
    plt.close()
    
    # Compute tail energy waste
    total_burst_energy = np.sum(data_activity) * (t[1] - t[0])
    total_power_energy = np.sum(power_state) * (t[1] - t[0])
    tail_waste_ratio = (total_power_energy - total_burst_energy) / total_power_energy
    
    return {
        "output_files": [f"{output_base}.pdf", f"{output_base}.png"],
        "computed_metrics": {
            "tail_waste_ratio": float(tail_waste_ratio),
            "tau_seconds": tau_tail
        },
        "validation_flags": {},
        "pass": True
    }