"""
Nature-standard figures for Campus Weather VAE paper.
"""
import os, sys, json
sys.path.insert(0, os.path.dirname(__file__))
import numpy as np
import matplotlib; matplotlib.use('Agg')
import matplotlib.pyplot as plt
from matplotlib.patches import FancyBboxPatch
from sklearn.decomposition import PCA

# Okabe-Ito colour-blind safe palette
C = {'blue':'#0072B2','orange':'#E69F00','green':'#009E73','red':'#D55E00',
     'purple':'#CC79A7','cyan':'#56B4E9','grey':'#999999','black':'#000000'}

plt.rcParams.update({
    'font.family': 'sans-serif', 'font.size': 7, 'axes.labelsize': 8,
    'axes.titlesize': 8, 'figure.dpi': 300, 'savefig.dpi': 300,
    'savefig.bbox': 'tight', 'axes.linewidth': 0.5, 'axes.spines.top': False,
    'axes.spines.right': False, 'axes.grid': False,
})

FIG = '/app/campus_weather/figures'
os.makedirs(FIG, exist_ok=True)


def panel_label(ax, label, x=-0.15, y=1.08):
    ax.text(x, y, label, transform=ax.transAxes, fontsize=10, fontweight='bold', va='bottom')


def fig1_campus_map(coords, cluster_labels):
    """Fig 1: Station map + discovered clusters."""
    fig, axes = plt.subplots(1, 2, figsize=(7.08, 3.2))
    
    # a: Station locations
    ax = axes[0]
    ax.scatter(coords[:, 1], coords[:, 0], c=C['blue'], s=25, edgecolors='white', linewidth=0.4, zorder=5)
    for i in range(len(coords)):
        ax.annotate(f'{i+1}', (coords[i,1], coords[i,0]), fontsize=3.5, ha='center', va='bottom',
                   xytext=(0,3), textcoords='offset points', color='#333')
    ax.set_xlabel('Longitude (°E)'); ax.set_ylabel('Latitude (°N)')
    ax.set_title('Station network (N=40)')
    panel_label(ax, 'a')
    
    # b: Clusters
    ax = axes[1]
    colours = [C['blue'], C['orange'], C['green'], C['red']]
    for c in range(max(cluster_labels)+1):
        mask = np.array(cluster_labels) == c
        ax.scatter(coords[mask, 1], coords[mask, 0], c=colours[c], s=35, 
                  edgecolors='black', linewidth=0.3, zorder=5, label=f'Zone {c+1}')
    ax.legend(fontsize=6, frameon=False, loc='lower right')
    ax.set_xlabel('Longitude (°E)'); ax.set_ylabel('Latitude (°N)')
    ax.set_title('Discovered microclimate zones (K=4)')
    panel_label(ax, 'b')
    
    plt.tight_layout(w_pad=1.0)
    plt.savefig(f'{FIG}/fig1_campus.pdf'); plt.savefig(f'{FIG}/fig1_campus.png'); plt.close()
    print('✓ Fig 1')


def fig2_reconstruction(results):
    """Fig 2: Reconstruction quality bar chart."""
    from train import VAR_NAMES, VAR_UNITS
    fig, ax = plt.subplots(1, 1, figsize=(3.54, 2.5))
    
    names = VAR_NAMES
    r2s = [results['reconstruction'][n]['R2'] for n in names]
    x = np.arange(len(names))
    bars = ax.bar(x, r2s, color=C['blue'], width=0.6, edgecolor='white', linewidth=0.3)
    ax.set_ylim(0.93, 1.001)
    ax.set_xticks(x); ax.set_xticklabels(names, rotation=45, ha='right', fontsize=6)
    ax.set_ylabel('R²')
    ax.set_title('Reconstruction quality (test set)')
    ax.axhline(0.99, color=C['grey'], linestyle='--', linewidth=0.5)
    
    for bar, val in zip(bars, r2s):
        ax.text(bar.get_x() + bar.get_width()/2, val + 0.001, f'{val:.4f}', 
                ha='center', va='bottom', fontsize=5)
    
    plt.tight_layout()
    plt.savefig(f'{FIG}/fig2_reconstruction.pdf'); plt.savefig(f'{FIG}/fig2_reconstruction.png'); plt.close()
    print('✓ Fig 2')


def fig3_spatial_interpolation(results):
    """Fig 3: Spatial interpolation — held-out station performance."""
    from train import VAR_NAMES
    fig, axes = plt.subplots(1, 2, figsize=(7.08, 2.8))
    
    si = results['spatial_interpolation']
    stations = [k for k in si if k.startswith('WS')]
    
    # a: AirTemp per station
    ax = axes[0]
    mae_vals = [si[s]['AirTemp']['MAE'] for s in stations]
    r2_vals = [si[s]['AirTemp']['R2'] for s in stations]
    x = np.arange(len(stations))
    ax.bar(x, mae_vals, color=C['orange'], width=0.6)
    ax.set_xticks(x); ax.set_xticklabels(stations, fontsize=6)
    ax.set_ylabel('MAE (°C)'); ax.set_title('Air temperature at held-out stations')
    panel_label(ax, 'a')
    
    # b: All variables average
    ax = axes[1]
    avg = si['average']
    vars_show = ['AirTemp', 'RelHum', 'AtmPress', 'WindSpeed']
    mae_avg = [avg[v]['MAE'] for v in vars_show]
    r2_avg = [avg[v]['R2'] for v in vars_show]
    x = np.arange(len(vars_show))
    bars = ax.bar(x - 0.15, mae_avg, 0.3, label='MAE', color=C['blue'])
    ax2 = ax.twinx()
    ax2.bar(x + 0.15, r2_avg, 0.3, label='R²', color=C['green'], alpha=0.7)
    ax.set_xticks(x); ax.set_xticklabels(vars_show, fontsize=6)
    ax.set_ylabel('MAE'); ax2.set_ylabel('R²')
    ax.set_title('Average across held-out stations')
    ax.legend(fontsize=5, loc='upper left', frameon=False)
    ax2.legend(fontsize=5, loc='upper right', frameon=False)
    panel_label(ax, 'b')
    
    plt.tight_layout(w_pad=1.0)
    plt.savefig(f'{FIG}/fig3_spatial.pdf'); plt.savefig(f'{FIG}/fig3_spatial.png'); plt.close()
    print('✓ Fig 3')


def fig4_forecasting(results):
    """Fig 4: Forecasting — embedding vs persistence vs climatology."""
    fc = results['temporal_forecasting']
    
    fig, axes = plt.subplots(1, 3, figsize=(7.08, 2.5))
    vars_show = ['AirTemp', 'RelHum', 'GlobalRad']
    units = ['°C', '%', 'W/m²']
    horizons = ['T+1', 'T+6', 'T+24']
    
    for ax, var, unit, pl in zip(axes, vars_show, units, ['a','b','c']):
        emb = [fc[h][var]['MAE_embedding'] for h in horizons]
        per = [fc[h][var]['MAE_persistence'] for h in horizons]
        clm = [fc[h][var]['MAE_climatology'] for h in horizons]
        
        x = np.arange(len(horizons))
        w = 0.25
        ax.bar(x - w, emb, w, label='Embedding', color=C['blue'])
        ax.bar(x, per, w, label='Persistence', color=C['orange'])
        ax.bar(x + w, clm, w, label='Climatology', color=C['grey'])
        
        ax.set_xticks(x); ax.set_xticklabels(horizons, fontsize=6)
        ax.set_ylabel(f'MAE ({unit})')
        ax.set_title(var)
        if pl == 'a': ax.legend(fontsize=5, frameon=False)
        panel_label(ax, pl)
    
    plt.tight_layout(w_pad=0.8)
    plt.savefig(f'{FIG}/fig4_forecasting.pdf'); plt.savefig(f'{FIG}/fig4_forecasting.png'); plt.close()
    print('✓ Fig 4')


def fig5_anomaly(results, data):
    """Fig 5: Anomaly detection — error timeseries + hour distribution."""
    ad = results['anomaly_detection']
    
    fig, axes = plt.subplots(1, 2, figsize=(7.08, 2.5))
    
    # a: Error timeseries (use saved npy)
    ax = axes[0]
    errors = np.load('/app/campus_weather/results/anomaly_errors.npy')
    # Downsample for plotting (daily mean)
    daily = errors.reshape(-1, 24).mean(axis=1)
    ax.plot(np.arange(len(daily)), daily, linewidth=0.5, color=C['blue'])
    threshold_daily = np.percentile(daily, 95)
    ax.axhline(threshold_daily, color=C['red'], linestyle='--', linewidth=0.8, label='95th percentile')
    ax.set_xlabel('Day of year'); ax.set_ylabel('Reconstruction error')
    ax.set_title('Campus-wide anomaly score')
    ax.legend(fontsize=5, frameon=False)
    panel_label(ax, 'a')
    
    # b: Hour distribution of anomalies
    ax = axes[1]
    hour_dist = np.array(ad['anomaly_hour_distribution'])
    ax.bar(np.arange(24), hour_dist, color=C['orange'], width=0.8)
    ax.set_xlabel('Hour of day'); ax.set_ylabel('Count')
    ax.set_title('Temporal distribution of anomalies')
    ax.set_xticks([0, 6, 12, 18])
    panel_label(ax, 'b')
    
    plt.tight_layout(w_pad=1.0)
    plt.savefig(f'{FIG}/fig5_anomaly.pdf'); plt.savefig(f'{FIG}/fig5_anomaly.png'); plt.close()
    print('✓ Fig 5')


def fig6_future(results):
    """Fig 6: 24h rolling forecast skill curve."""
    fp = results['future_prediction']['per_hour']
    
    fig, axes = plt.subplots(1, 3, figsize=(7.08, 2.5))
    vars_show = ['AirTemp', 'RelHum', 'GlobalRad']
    units = ['°C', '%', 'W/m²']
    
    hours = list(range(1, 25))
    
    for ax, var, unit, pl in zip(axes, vars_show, units, ['a','b','c']):
        mae_emb = [fp[f'h+{h}'][var]['MAE_embedding'] for h in hours]
        mae_per = [fp[f'h+{h}'][var]['MAE_persistence'] for h in hours]
        mae_clm = [fp[f'h+{h}'][var]['MAE_climatology'] for h in hours]
        
        ax.plot(hours, mae_emb, color=C['blue'], linewidth=1.2, label='Embedding')
        ax.plot(hours, mae_per, color=C['orange'], linewidth=1.2, linestyle='--', label='Persistence')
        ax.plot(hours, mae_clm, color=C['grey'], linewidth=1.2, linestyle=':', label='Climatology')
        
        ax.set_xlabel('Forecast horizon (h)')
        ax.set_ylabel(f'MAE ({unit})')
        ax.set_title(var)
        ax.set_xticks([1, 6, 12, 18, 24])
        if pl == 'a': ax.legend(fontsize=5, frameon=False)
        panel_label(ax, pl)
    
    plt.tight_layout(w_pad=0.8)
    plt.savefig(f'{FIG}/fig6_future.pdf'); plt.savefig(f'{FIG}/fig6_future.png'); plt.close()
    print('✓ Fig 6')


if __name__ == '__main__':
    results = json.load(open('/app/campus_weather/results/all_results.json'))
    npz = np.load('/app/campus_weather/results/checkpoints/embeddings.npz', allow_pickle=True)
    data, coords = npz['data'], npz['coords']
    
    cluster_labels = results['clustering']['K=4']['labels']
    
    fig1_campus_map(coords, cluster_labels)
    fig2_reconstruction(results)
    fig3_spatial_interpolation(results)
    fig4_forecasting(results)
    fig5_anomaly(results, data)
    fig6_future(results)
    
    print(f'\nAll figures saved to {FIG}/')