"""
EPW Weather File Generation from VAE Embeddings
================================================
Generates EnergyPlus Weather (EPW) files for building energy simulation:

1. BASELINE: Campus-mean weather from imputed observations (2025)
2. HEATWAVE: +2°C sustained warming via latent space manipulation
3. UHI INTENSIFICATION: +2°C warming with reduced wind (urban canyon effect)

The VAE's continuous latent space enables generation of physically coherent
extreme scenarios where all six weather variables shift together consistently —
something classical interpolation methods (IDW, kriging) cannot do.

Run: python generate_epw.py
Outputs: epw/NUS_baseline_2025.epw, epw/NUS_heatwave.epw, epw/NUS_uhi.epw
"""

import sys, os, json, warnings
sys.path.insert(0, os.path.dirname(__file__))
warnings.filterwarnings('ignore')

import numpy as np
import pandas as pd
import pvlib
import torch
from sklearn.linear_model import LinearRegression
from ladybug.epw import EPW
from ladybug.location import Location

from model import WeatherVAE
from train import load_nus40, VAR_NAMES, VAR_UNITS

# ── Paths ────────────────────────────────────────────────────────────────
BASE = '/app/campus_weather'
DATA_DIR = f'{BASE}/imputed'
RESULTS = f'{BASE}/results'
EPW_DIR = f'{BASE}/epw'
FIG_DIR = f'{BASE}/figures'

# ── Campus location ──────────────────────────────────────────────────────
CAMPUS_LAT = 1.2992
CAMPUS_LON = 103.7764
CAMPUS_ALT = 16  # metres ASL
TZ_OFFSET = 8    # UTC+8
TZ_STR = 'Asia/Singapore'
YEAR = 2025


# ═══════════════════════════════════════════════════════════════════════
# DERIVED FIELD COMPUTATIONS
# ═══════════════════════════════════════════════════════════════════════

def dewpoint_magnus(T_C, RH_pct):
    """Dew point temperature via Magnus formula. T in °C, RH in %."""
    RH_safe = np.clip(RH_pct, 1, 100)
    a, b = 17.27, 237.7
    gamma = a * T_C / (b + T_C) + np.log(RH_safe / 100.0)
    return b * gamma / (a - gamma)


def horizontal_infrared_radiation(T_C, RH_pct):
    """
    Horizontal infrared radiation intensity (W/m²).
    Martin-Berdahl model using dew point temperature.
    """
    sigma = 5.67e-8
    T_K = T_C + 273.15
    Tdp_K = dewpoint_magnus(T_C, RH_pct) + 273.15
    eps_sky = np.clip(0.787 + 0.764 * np.log(Tdp_K / 273.16), 0.6, 1.0)
    return eps_sky * sigma * T_K**4


def compute_solar_derived(ghr, year=YEAR):
    """
    From Global Horizontal Radiation (W/m²), compute:
    - DNI (Direct Normal Irradiance)
    - DHI (Diffuse Horizontal Irradiance)
    - Extraterrestrial radiation
    - Sky cover (tenths)
    
    All returned in Wh/m² (numerically equal to W/m² for hourly averages).
    """
    times = pd.date_range(f'{year}-01-01', periods=8760, freq='1h', tz=TZ_STR)
    solpos = pvlib.solarposition.get_solarposition(times, CAMPUS_LAT, CAMPUS_LON, CAMPUS_ALT)
    zenith = solpos['apparent_zenith'].values

    ghr_clean = np.clip(ghr, 0, 1400)
    # Zero out nighttime
    ghr_clean[zenith > 90] = 0

    decomp = pvlib.irradiance.erbs(ghr_clean, zenith, times)
    dni = np.nan_to_num(decomp['dni'].values, nan=0.0).clip(0, 1200)
    dhi = np.nan_to_num(decomp['dhi'].values, nan=0.0).clip(0, 800)

    # Extraterrestrial
    etrn = pvlib.irradiance.get_extra_radiation(times).values  # Direct normal
    etr = (etrn * np.cos(np.radians(zenith))).clip(0)  # Horizontal

    # Sky cover from clearness index
    kt = np.where(etr > 10, ghr_clean / etr, 0).clip(0, 1)
    sky_cover = ((1 - kt) * 10).clip(0, 10).astype(int)

    return {
        'dni': dni, 'dhi': dhi,
        'etrn': etrn, 'etr': etr,
        'sky_cover': sky_cover, 'zenith': zenith,
    }


def physical_bounds_clip(weather):
    """Clip decoded weather to physical bounds."""
    weather[:, 0] = np.clip(weather[:, 0], 0, 20)       # WindSpeed: 0-20 m/s
    weather[:, 1] = np.clip(weather[:, 1], 0, 360)      # WindDir: 0-360°
    weather[:, 2] = np.clip(weather[:, 2], 15, 50)       # AirTemp: 15-50°C
    weather[:, 3] = np.clip(weather[:, 3], 10, 100)      # RelHum: 10-100%
    weather[:, 4] = np.clip(weather[:, 4], 990, 1020)    # AtmPress: 990-1020 hPa
    weather[:, 5] = np.clip(weather[:, 5], 0, 1400)      # GlobalRad: 0-1400 W/m²
    return weather


# ═══════════════════════════════════════════════════════════════════════
# EPW FILE GENERATION
# ═══════════════════════════════════════════════════════════════════════

def weather_to_epw(weather, scenario_name, description, output_path, station_id='NUS_VAE'):
    """
    Convert a (8760, 6) weather array to a valid EPW file.
    
    weather columns: [WindSpeed, WindDir, AirTemp, RelHum, AtmPress, GlobalRad]
    """
    ws  = weather[:, 0]  # m/s
    wd  = weather[:, 1]  # degrees
    ta  = weather[:, 2]  # °C
    rh  = weather[:, 3]  # %
    pa  = weather[:, 4] * 100  # hPa → Pa (EPW uses Pa)
    ghr = weather[:, 5]  # W/m² = Wh/m² for hourly

    # Derived fields
    dp = dewpoint_magnus(ta, rh)
    hiri = horizontal_infrared_radiation(ta, rh)
    solar = compute_solar_derived(ghr)

    # Build EPW
    epw = EPW.from_missing_values()
    epw.location = Location(
        city='NUS_Campus', state='', country='SGP',
        source=f'VAE_{scenario_name}', station_id=station_id,
        latitude=CAMPUS_LAT, longitude=CAMPUS_LON,
        time_zone=TZ_OFFSET, elevation=CAMPUS_ALT
    )
    epw.comments_1 = f'VAE-generated EPW: {scenario_name}'
    epw.comments_2 = description

    # Assign fields (all 8760-length lists)
    epw.years.values = [YEAR] * 8760
    epw.dry_bulb_temperature.values = ta.tolist()
    epw.dew_point_temperature.values = dp.tolist()
    epw.relative_humidity.values = rh.tolist()
    epw.atmospheric_station_pressure.values = pa.tolist()
    epw.wind_speed.values = ws.tolist()
    epw.wind_direction.values = wd.tolist()
    epw.global_horizontal_radiation.values = ghr.tolist()
    epw.direct_normal_radiation.values = solar['dni'].tolist()
    epw.diffuse_horizontal_radiation.values = solar['dhi'].tolist()
    epw.horizontal_infrared_radiation_intensity.values = hiri.tolist()
    epw.extraterrestrial_horizontal_radiation.values = solar['etr'].tolist()
    epw.extraterrestrial_direct_normal_radiation.values = solar['etrn'].tolist()
    epw.total_sky_cover.values = solar['sky_cover'].tolist()
    epw.opaque_sky_cover.values = solar['sky_cover'].tolist()

    # Singapore-specific fixed fields
    epw.snow_depth.values = [0] * 8760
    epw.days_since_last_snowfall.values = [99] * 8760
    epw.albedo.values = [0.15] * 8760  # typical urban

    os.makedirs(os.path.dirname(output_path), exist_ok=True)
    epw.save(output_path)
    return epw


# ═══════════════════════════════════════════════════════════════════════
# SCENARIO GENERATION
# ═══════════════════════════════════════════════════════════════════════

def generate_scenarios():
    """Generate all three EPW scenarios."""
    os.makedirs(EPW_DIR, exist_ok=True)

    # Load data and model
    print("Loading data and model...")
    data, coords, datetimes = load_nus40(DATA_DIR)
    T, N, V = data.shape

    ckpt = torch.load(f'{RESULTS}/checkpoints/best.pt', map_location='cpu', weights_only=False)
    model = WeatherVAE(**ckpt['config'])
    model.load_state_dict(ckpt['model'])
    model.set_normalisation(ckpt['mean'], ckpt['std'])
    model.eval()

    emb = np.load(f'{RESULTS}/checkpoints/embeddings.npz')['embeddings']
    campus_emb = emb.mean(axis=1)   # (8760, 6) — campus-mean embedding per hour
    campus_data = data.mean(axis=1)  # (8760, 6) — campus-mean weather per hour

    # ── Find latent directions via linear regression ─────────────────
    print("Computing latent directions...")
    directions = {}
    for v, name in enumerate(VAR_NAMES):
        reg = LinearRegression()
        reg.fit(campus_emb, campus_data[:, v])
        w = reg.coef_
        directions[name] = w / np.linalg.norm(w)

    # ── Decode baseline to get the VAE's reconstruction ──────────────
    with torch.no_grad():
        baseline_decoded = model.decode(
            torch.from_numpy(campus_emb.astype(np.float32))
        ).numpy()
    baseline_decoded = physical_bounds_clip(baseline_decoded)

    # ═════════════════════════════════════════════════════════════════
    # SCENARIO 1: BASELINE
    # Campus-mean imputed observations — the actual 2025 weather
    # ═════════════════════════════════════════════════════════════════
    print("\n" + "=" * 60)
    print("SCENARIO 1: BASELINE (imputed observations)")
    print("=" * 60)

    baseline_weather = campus_data.copy()
    epw_baseline = weather_to_epw(
        baseline_weather,
        scenario_name='Baseline_2025',
        description='Campus-mean of 40 NUS stations, imputed observations, Jan-Dec 2025',
        output_path=f'{EPW_DIR}/NUS_baseline_2025.epw',
        station_id='NUS_BAS_2025',
    )
    print_scenario_stats('Baseline', baseline_weather)

    # ═════════════════════════════════════════════════════════════════
    # SCENARIO 2: HEATWAVE (+2°C)
    # Shift the entire year's embeddings along the temperature
    # direction in latent space, then decode. This produces a
    # physically coherent heatwave where humidity drops, radiation
    # increases, and pressure decreases — all simultaneously.
    # ═════════════════════════════════════════════════════════════════
    print("\n" + "=" * 60)
    print("SCENARIO 2: HEATWAVE (+2°C via latent shift)")
    print("=" * 60)

    # Calibrate: scale=0.75 in AirTemp direction → ~+2°C
    # Verified empirically above
    heat_scale = 0.75
    z_heat = campus_emb + heat_scale * directions['AirTemp']

    with torch.no_grad():
        heat_decoded = model.decode(
            torch.from_numpy(z_heat.astype(np.float32))
        ).numpy()
    heat_decoded = physical_bounds_clip(heat_decoded)

    # Verify the shift magnitude
    dt_mean = heat_decoded[:, 2].mean() - baseline_decoded[:, 2].mean()
    print(f"  Achieved mean ΔT: {dt_mean:+.2f}°C (target: +2°C)")

    # If not close enough, adjust scale
    if abs(dt_mean - 2.0) > 0.3:
        # Binary search for correct scale
        lo, hi = 0.1, 3.0
        for _ in range(20):
            mid = (lo + hi) / 2
            z_test = campus_emb + mid * directions['AirTemp']
            with torch.no_grad():
                dec_test = model.decode(torch.from_numpy(z_test.astype(np.float32))).numpy()
            dt_test = dec_test[:, 2].mean() - baseline_decoded[:, 2].mean()
            if dt_test < 2.0:
                lo = mid
            else:
                hi = mid
        heat_scale = (lo + hi) / 2
        z_heat = campus_emb + heat_scale * directions['AirTemp']
        with torch.no_grad():
            heat_decoded = model.decode(torch.from_numpy(z_heat.astype(np.float32))).numpy()
        heat_decoded = physical_bounds_clip(heat_decoded)
        dt_mean = heat_decoded[:, 2].mean() - baseline_decoded[:, 2].mean()
        print(f"  Recalibrated: scale={heat_scale:.3f}, ΔT={dt_mean:+.2f}°C")

    epw_heat = weather_to_epw(
        heat_decoded,
        scenario_name='Heatwave_plus2C',
        description=f'Heatwave scenario: +{dt_mean:.1f}C sustained warming via VAE latent shift (scale={heat_scale:.3f})',
        output_path=f'{EPW_DIR}/NUS_heatwave.epw',
        station_id='NUS_HW_2025',
    )
    print_scenario_stats('Heatwave', heat_decoded)

    # ═════════════════════════════════════════════════════════════════
    # SCENARIO 3: URBAN HEAT ISLAND INTENSIFICATION
    # Combined shift: temperature up AND wind speed down.
    # This represents the effect of increased building density
    # reducing ventilation corridors while trapping more heat.
    # Physically distinct from a heatwave: UHI warming is strongest
    # at night (reduced nocturnal cooling), whereas heatwaves
    # affect daytime peaks.
    # ═════════════════════════════════════════════════════════════════
    print("\n" + "=" * 60)
    print("SCENARIO 3: UHI INTENSIFICATION (+2°C, -30% wind)")
    print("=" * 60)

    # Combine: push temperature up + push wind down
    uhi_dir = directions['AirTemp'] - 0.5 * directions['WindSpeed']
    uhi_dir = uhi_dir / np.linalg.norm(uhi_dir)

    # Calibrate for +2°C
    lo, hi = 0.1, 5.0
    for _ in range(20):
        mid = (lo + hi) / 2
        z_test = campus_emb + mid * uhi_dir
        with torch.no_grad():
            dec_test = model.decode(torch.from_numpy(z_test.astype(np.float32))).numpy()
        dt_test = dec_test[:, 2].mean() - baseline_decoded[:, 2].mean()
        if dt_test < 2.0:
            lo = mid
        else:
            hi = mid
    uhi_scale = (lo + hi) / 2

    z_uhi = campus_emb + uhi_scale * uhi_dir
    with torch.no_grad():
        uhi_decoded = model.decode(torch.from_numpy(z_uhi.astype(np.float32))).numpy()
    uhi_decoded = physical_bounds_clip(uhi_decoded)

    dt_uhi = uhi_decoded[:, 2].mean() - baseline_decoded[:, 2].mean()
    dws_uhi = uhi_decoded[:, 0].mean() - baseline_decoded[:, 0].mean()
    ws_pct = dws_uhi / baseline_decoded[:, 0].mean() * 100
    print(f"  Achieved: ΔT={dt_uhi:+.2f}°C, ΔWS={dws_uhi:+.3f} m/s ({ws_pct:+.0f}%)")

    epw_uhi = weather_to_epw(
        uhi_decoded,
        scenario_name='UHI_Intensification',
        description=f'Urban heat island intensification: +{dt_uhi:.1f}C warming, {ws_pct:.0f}% wind reduction via VAE latent shift (scale={uhi_scale:.3f})',
        output_path=f'{EPW_DIR}/NUS_uhi.epw',
        station_id='NUS_UHI_2025',
    )
    print_scenario_stats('UHI', uhi_decoded)

    # ═════════════════════════════════════════════════════════════════
    # COMPARISON STATISTICS
    # ═════════════════════════════════════════════════════════════════
    print("\n" + "=" * 60)
    print("SCENARIO COMPARISON")
    print("=" * 60)

    scenarios = {
        'Baseline (imputed)': baseline_weather,
        'Baseline (VAE decoded)': baseline_decoded,
        'Heatwave (+2°C)': heat_decoded,
        'UHI Intensification': uhi_decoded,
    }

    print(f"\n{'Scenario':<25s} | {'AirTemp':>8s} {'RelHum':>8s} {'WindSpd':>8s} {'GloRad':>8s} {'Press':>8s}")
    print("-" * 75)
    for name, w in scenarios.items():
        print(f"{name:<25s} | {w[:, 2].mean():>7.1f}° {w[:, 3].mean():>7.1f}% "
              f"{w[:, 0].mean():>7.2f}m {w[:, 5].mean():>7.1f}W {w[:, 4].mean():>7.1f}h")

    # Deltas from baseline
    print(f"\n{'Scenario':<25s} | {'ΔTemp':>8s} {'ΔRH':>8s} {'ΔWind':>8s} {'ΔRad':>8s} {'ΔPress':>8s}")
    print("-" * 75)
    for name, w in [('Heatwave', heat_decoded), ('UHI', uhi_decoded)]:
        dt = w[:, 2].mean() - baseline_decoded[:, 2].mean()
        dr = w[:, 3].mean() - baseline_decoded[:, 3].mean()
        dw = w[:, 0].mean() - baseline_decoded[:, 0].mean()
        dg = w[:, 5].mean() - baseline_decoded[:, 5].mean()
        dp = w[:, 4].mean() - baseline_decoded[:, 4].mean()
        print(f"{name:<25s} | {dt:>+7.2f}° {dr:>+7.1f}% {dw:>+7.3f}m {dg:>+7.1f}W {dp:>+7.2f}h")

    # Diurnal profiles
    hours = np.arange(8760) % 24
    print(f"\nDiurnal Temperature Profile:")
    print(f"{'Hour':>6s} | {'Baseline':>10s} {'Heatwave':>10s} {'UHI':>10s} | {'ΔHW':>6s} {'ΔUHI':>6s}")
    print("-" * 60)
    for h in range(0, 24, 3):
        mask = hours == h
        b = baseline_decoded[mask, 2].mean()
        hw = heat_decoded[mask, 2].mean()
        uhi = uhi_decoded[mask, 2].mean()
        print(f"{h:>6d} | {b:>10.2f} {hw:>10.2f} {uhi:>10.2f} | {hw-b:>+5.2f} {uhi-b:>+5.2f}")

    # Nighttime vs daytime warming (UHI signature check)
    day_mask = (hours >= 8) & (hours <= 18)
    night_mask = ~day_mask

    print(f"\nDay vs Night warming (UHI should warm more at night):")
    for name, w in [('Heatwave', heat_decoded), ('UHI', uhi_decoded)]:
        day_dt = w[day_mask, 2].mean() - baseline_decoded[day_mask, 2].mean()
        night_dt = w[night_mask, 2].mean() - baseline_decoded[night_mask, 2].mean()
        print(f"  {name}: Day ΔT={day_dt:+.2f}°C, Night ΔT={night_dt:+.2f}°C, "
              f"Night-Day={night_dt - day_dt:+.2f}°C")

    # Save scenario metadata
    results = {
        'scenarios': {},
        'latent_directions': {name: d.tolist() for name, d in directions.items()},
    }
    for sname, w, scale in [
        ('baseline_imputed', baseline_weather, None),
        ('baseline_vae', baseline_decoded, None),
        ('heatwave', heat_decoded, heat_scale),
        ('uhi', uhi_decoded, uhi_scale),
    ]:
        stats = {}
        for v, vname in enumerate(VAR_NAMES):
            stats[vname] = {
                'mean': round(float(w[:, v].mean()), 3),
                'std': round(float(w[:, v].std()), 3),
                'min': round(float(w[:, v].min()), 3),
                'max': round(float(w[:, v].max()), 3),
            }
        if scale is not None:
            stats['latent_scale'] = round(float(scale), 4)
        results['scenarios'][sname] = stats

    with open(f'{EPW_DIR}/scenario_results.json', 'w') as f:
        json.dump(results, f, indent=2)
    print(f"\nResults saved to {EPW_DIR}/scenario_results.json")

    # Validate EPW files
    print("\n" + "=" * 60)
    print("EPW VALIDATION")
    print("=" * 60)
    validate_epw_files()

    return results


def print_scenario_stats(name, weather):
    """Print summary statistics for a scenario."""
    print(f"\n  {name} annual statistics:")
    for v, (vname, unit) in enumerate(zip(VAR_NAMES, VAR_UNITS)):
        col = weather[:, v]
        print(f"    {vname:>12s}: mean={col.mean():.2f} std={col.std():.2f} "
              f"min={col.min():.2f} max={col.max():.2f} {unit}")


def validate_epw_files():
    """Validate generated EPW files by reading them back."""
    for fname in ['NUS_baseline_2025.epw', 'NUS_heatwave.epw', 'NUS_uhi.epw']:
        path = f'{EPW_DIR}/{fname}'
        if not os.path.exists(path):
            print(f"  ✗ {fname}: NOT FOUND")
            continue

        try:
            epw = EPW(path)
            ta = np.array(epw.dry_bulb_temperature.values)
            rh = np.array(epw.relative_humidity.values)
            ws = np.array(epw.wind_speed.values)
            ghr = np.array(epw.global_horizontal_radiation.values)
            dni = np.array(epw.direct_normal_radiation.values)
            dhi = np.array(epw.diffuse_horizontal_radiation.values)

            n_hours = len(ta)
            ta_range = f"{ta.min():.1f}–{ta.max():.1f}°C"
            rh_range = f"{rh.min():.0f}–{rh.max():.0f}%"

            # Physical checks
            issues = []
            if (ghr < -1).any(): issues.append("GHR < 0")
            if (dni < -1).any(): issues.append("DNI < 0")
            if (ta < -50).any() or (ta > 60).any(): issues.append("T out of range")
            if (rh < 0).any() or (rh > 101).any(): issues.append("RH out of range")
            if (ws < 0).any(): issues.append("WS < 0")

            # Check GHR = 0 at night (approximately)
            night_ghr = ghr[np.arange(8760) % 24 < 6]  # midnight to 6am
            if night_ghr.max() > 10:
                issues.append(f"GHR at night: max={night_ghr.max():.1f}")

            # Check DNI + DHI consistency (DNI*cos(z) + DHI ≈ GHR)
            # Only check during daytime
            day_mask = ghr > 50
            if day_mask.sum() > 100:
                ghr_check = ghr[day_mask]
                dni_check = dni[day_mask]
                dhi_check = dhi[day_mask]
                ratio = (dni_check + dhi_check) / np.maximum(ghr_check, 1)
                # DNI here is direct normal, not horizontal. Skip strict check.

            status = '✓' if not issues else f"⚠ {', '.join(issues)}"
            print(f"  {status} {fname}: {n_hours} hours, T={ta_range}, RH={rh_range}, "
                  f"mean GHR={ghr.mean():.0f} W/m²")

        except Exception as e:
            print(f"  ✗ {fname}: FAILED TO READ — {e}")


def make_scenario_figures():
    """Generate comparison figures for the three scenarios."""
    import matplotlib
    matplotlib.use('Agg')
    import matplotlib.pyplot as plt
    import seaborn as sns

    os.makedirs(FIG_DIR, exist_ok=True)
    C = ['#264653', '#2a9d8f', '#e9c46a', '#e76f51']
    sns.set_theme(style='whitegrid', font_scale=1.05)

    # Load the three EPW files
    epw_files = {
        'Baseline': f'{EPW_DIR}/NUS_baseline_2025.epw',
        'Heatwave (+2°C)': f'{EPW_DIR}/NUS_heatwave.epw',
        'UHI Intensification': f'{EPW_DIR}/NUS_uhi.epw',
    }

    all_data = {}
    for name, path in epw_files.items():
        epw = EPW(path)
        all_data[name] = {
            'AirTemp': np.array(epw.dry_bulb_temperature.values),
            'RelHum': np.array(epw.relative_humidity.values),
            'WindSpeed': np.array(epw.wind_speed.values),
            'GlobalRad': np.array(epw.global_horizontal_radiation.values),
            'DewPoint': np.array(epw.dew_point_temperature.values),
            'Pressure': np.array(epw.atmospheric_station_pressure.values) / 100,  # Pa → hPa
        }

    hours = np.arange(8760) % 24
    months = np.array([(pd.Timestamp(f'{YEAR}-01-01') + pd.Timedelta(hours=h)).month for h in range(8760)])

    # ── Fig 14: Diurnal profiles — 4 panels ──────────────────────────
    fig, axes = plt.subplots(2, 2, figsize=(12, 9))
    colors = [C[0], C[3], C[1]]

    for ax, var, ylabel in zip(axes.flat,
            ['AirTemp', 'RelHum', 'WindSpeed', 'GlobalRad'],
            ['Temperature (°C)', 'Relative Humidity (%)', 'Wind Speed (m/s)', 'Global Radiation (W/m²)']):
        for (name, d), color in zip(all_data.items(), colors):
            diurnal = [d[var][hours == h].mean() for h in range(24)]
            ax.plot(range(24), diurnal, color=color, lw=2, label=name)
        ax.set_xlabel('Hour of day')
        ax.set_ylabel(ylabel)
        ax.set_xlim(0, 23)
        ax.set_xticks([0, 3, 6, 9, 12, 15, 18, 21])
        ax.legend(fontsize=8)

    axes[0, 0].set_title('(a) Air temperature')
    axes[0, 1].set_title('(b) Relative humidity')
    axes[1, 0].set_title('(c) Wind speed')
    axes[1, 1].set_title('(d) Global solar radiation')

    plt.suptitle('Diurnal Profiles: Baseline vs Extreme Scenarios', fontsize=13, y=1.01)
    plt.tight_layout()
    for ext in ['png', 'pdf']:
        fig.savefig(f'{FIG_DIR}/fig14_epw_diurnal.{ext}', dpi=300, bbox_inches='tight')
    plt.close()
    print("  Saved fig14_epw_diurnal")

    # ── Fig 15: Monthly temperature box plots ────────────────────────
    fig, axes = plt.subplots(1, 3, figsize=(15, 4.5), sharey=True)

    for ax, (name, d), color in zip(axes, all_data.items(), colors):
        month_data = [d['AirTemp'][months == m] for m in range(1, 13)]
        bp = ax.boxplot(month_data, patch_artist=True, widths=0.6,
                       medianprops=dict(color='black', lw=1.5))
        for patch in bp['boxes']:
            patch.set_facecolor(color)
            patch.set_alpha(0.6)
        ax.set_xlabel('Month')
        ax.set_title(name)
        ax.set_xticklabels(['J','F','M','A','M','J','J','A','S','O','N','D'])

    axes[0].set_ylabel('Air Temperature (°C)')
    plt.suptitle('Monthly Temperature Distributions', fontsize=13, y=1.02)
    plt.tight_layout()
    for ext in ['png', 'pdf']:
        fig.savefig(f'{FIG_DIR}/fig15_epw_monthly.{ext}', dpi=300, bbox_inches='tight')
    plt.close()
    print("  Saved fig15_epw_monthly")

    # ── Fig 16: Psychrometric-style scatter (T vs RH) ────────────────
    fig, ax = plt.subplots(figsize=(8, 6))

    rng = np.random.RandomState(42)
    n_sample = 2000
    for (name, d), color, marker in zip(all_data.items(), colors, ['o', '^', 's']):
        idx = rng.choice(8760, n_sample, replace=False)
        ax.scatter(d['AirTemp'][idx], d['RelHum'][idx],
                  c=color, alpha=0.15, s=10, marker=marker, label=name, rasterized=True)

    ax.set_xlabel('Air Temperature (°C)')
    ax.set_ylabel('Relative Humidity (%)')
    ax.set_title('Temperature–Humidity State Space')
    ax.legend(markerscale=3, fontsize=10)
    ax.set_xlim(20, 42)
    ax.set_ylim(25, 100)

    plt.tight_layout()
    for ext in ['png', 'pdf']:
        fig.savefig(f'{FIG_DIR}/fig16_epw_psychrometric.{ext}', dpi=300, bbox_inches='tight')
    plt.close()
    print("  Saved fig16_epw_psychrometric")


# ═══════════════════════════════════════════════════════════════════════

if __name__ == '__main__':
    results = generate_scenarios()
    print("\n" + "=" * 60)
    print("GENERATING FIGURES...")
    print("=" * 60)
    make_scenario_figures()
    print("\nDone.")