code/thermal_comfort.py

"""
Evaluation 6: Thermal comfort (UTCI) from weather embeddings.

Computes UTCI at all 40 stations using:
  - pvlib for solar geometry + irradiance decomposition (Erbs et al. 1982)
  - pythermalcomfort for solar_gain (MRT delta) and UTCI (Bröde et al. 2012)

Then tests whether the VAE embedding can spatially interpolate UTCI at
held-out stations — the practical task of campus comfort mapping from
sparse measurements.

Run: python thermal_comfort.py
Outputs: results/thermal_comfort.json, results/utci_all.npz, figures/fig9-11
"""

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

import numpy as np
import pandas as pd
import pvlib
from pythermalcomfort.models import utci as utci_fn, solar_gain as solar_gain_fn
from sklearn.neighbors import NearestNeighbors
from sklearn.linear_model import Ridge
from sklearn.metrics import r2_score, mean_absolute_error, mean_squared_error

from train import load_nus40, VAR_NAMES, VAR_COLS

# ── Paths ────────────────────────────────────────────────────────────────────
DATA_DIR = '/app/campus_weather/imputed'
RESULTS_DIR = '/app/campus_weather/results'
FIG_DIR = '/app/campus_weather/figures'
CKPT_DIR = os.path.join(RESULTS_DIR, 'checkpoints')

# ── Constants ────────────────────────────────────────────────────────────────
TZ = 'Asia/Singapore'  # UTC+8, no DST
CAMPUS_LAT = 1.2992    # centroid of the 40 stations
CAMPUS_LON = 103.7764
CAMPUS_ALT = 15        # metres ASL

# Measurement assumptions
Z_MEAS = 2.0      # anemometer height [m] — typical for campus AWS
Z0 = 0.1          # roughness length [m] — suburban/campus
F_SVV = 0.85       # sky view factor — mostly open campus with some buildings
F_BES = 0.9        # body exposure to sun

# UTCI stress category thresholds
UTCI_CATS = [
    (-np.inf, 9,   'no thermal stress'),
    (9,       26,  'no thermal stress'),
    (26,      32,  'moderate heat stress'),
    (32,      38,  'strong heat stress'),
    (38,      46,  'very strong heat stress'),
    (46,      np.inf, 'extreme heat stress'),
]


def wind_height_correction(v_meas, z_meas=Z_MEAS, z_target=10.0, z0=Z0):
    """Log-law wind profile: adjust measurement height to 10 m (UTCI standard)."""
    ratio = np.log(z_target / z0) / np.log(z_meas / z0)
    return np.clip(v_meas * ratio, 0.5, None)  # UTCI valid for v >= 0.5 m/s


def compute_solar_position(datetimes):
    """
    Solar position for the campus centroid.
    All 40 stations span ~2 km — solar geometry difference is negligible.
    Returns solar altitude [deg] and azimuth [deg].
    """
    loc = pvlib.location.Location(CAMPUS_LAT, CAMPUS_LON, tz=TZ, altitude=CAMPUS_ALT)
    times = pd.DatetimeIndex(datetimes)
    if times.tz is None:
        times = times.tz_localize(TZ)
    sol_pos = loc.get_solarposition(times)
    return sol_pos['apparent_elevation'].values, sol_pos['azimuth'].values


def decompose_irradiance(ghi, solar_zenith, datetimes):
    """
    Split GHI into DNI + DHI using Erbs et al. (1982) model.
    Returns DNI [W/m²] and DHI [W/m²].
    """
    times = pd.DatetimeIndex(datetimes)
    if times.tz is None:
        times = times.tz_localize(TZ)

    ghi_clean = np.clip(ghi, 0, 1400)
    decomp = pvlib.irradiance.erbs(
        ghi=ghi_clean,
        zenith=solar_zenith,
        datetime_or_doy=times
    )
    dni = np.nan_to_num(decomp['dni'].values, nan=0.0).clip(0)
    dhi = np.nan_to_num(decomp['dhi'].values, nan=0.0).clip(0)
    return dni, dhi


def compute_delta_mrt(solar_alt, dni):
    """
    Compute MRT increment from solar radiation using pythermalcomfort.solar_gain.
    Vectorised over time. Returns delta_MRT [°C].
    """
    T = len(solar_alt)
    delta_mrt = np.zeros(T)

    # Only compute for valid daytime hours
    valid = (solar_alt > 3.0) & (dni > 5.0)
    idx = np.where(valid)[0]

    if len(idx) == 0:
        return delta_mrt

    # pythermalcomfort accepts lists — batch all valid hours
    sg = solar_gain_fn(
        sol_altitude=solar_alt[idx].tolist(),
        sharp=[135.0] * len(idx),        # random body orientation
        sol_radiation_dir=dni[idx].tolist(),
        sol_transmittance=[1.0] * len(idx),
        f_svv=[F_SVV] * len(idx),
        f_bes=[F_BES] * len(idx),
        asw=0.7,
        posture='standing',
        floor_reflectance=0.6,
        round_output=False,
    )
    delta_mrt[idx] = np.asarray(sg.delta_mrt)
    return delta_mrt


def compute_utci_station(air_temp, rel_hum, wind_speed, delta_mrt):
    """
    Compute UTCI for one station (all hours). Returns UTCI [°C] array and stress categories.
    """
    T = len(air_temp)
    mrt = air_temp + delta_mrt  # MRT = Ta + delta from solar gain
    v10 = wind_height_correction(wind_speed)

    # Clip inputs to UTCI model validity range
    ta = np.clip(air_temp, -50, 50)
    tr = np.clip(mrt, -30, 70)
    v = np.clip(v10, 0.5, 17.0)
    rh = np.clip(rel_hum, 0, 100)

    result = utci_fn(
        tdb=ta.tolist(),
        tr=tr.tolist(),
        v=v.tolist(),
        rh=rh.tolist(),
        units='SI',
        limit_inputs=False,
        round_output=False,
    )
    utci_vals = np.asarray(result.utci)
    categories = np.asarray(result.stress_category)
    return utci_vals, categories, mrt


def compute_all_stations(data, datetimes):
    """
    Compute UTCI for all 40 stations.
    data: (T, N, V) array — [WindSpeed, WindDir, AirTemp, RelHum, AtmPress, GlobalRad]
    Returns: utci (T, N), mrt (T, N), categories (T, N)
    """
    T, N, V = data.shape

    # Solar position (same for all stations — 2 km campus)
    solar_alt, solar_az = compute_solar_position(datetimes)
    solar_zenith = 90.0 - solar_alt

    utci_all = np.zeros((T, N))
    mrt_all = np.zeros((T, N))
    cat_all = np.empty((T, N), dtype=object)

    for s in range(N):
        print(f'  Station {s+1:02d}/40', end='\r')
        ghi = data[:, s, 5]  # GlobalRad
        ta = data[:, s, 2]   # AirTemp
        rh = data[:, s, 3]   # RelHum
        ws = data[:, s, 0]   # WindSpeed

        # Decompose GHI → DNI
        dni, dhi = decompose_irradiance(ghi, solar_zenith, datetimes)

        # MRT delta from solar gain
        dmrt = compute_delta_mrt(solar_alt, dni)

        # UTCI
        utci_vals, cats, mrt = compute_utci_station(ta, rh, ws, dmrt)
        utci_all[:, s] = utci_vals
        mrt_all[:, s] = mrt
        cat_all[:, s] = cats

    print(f'  Computed UTCI for {N} stations × {T} hours')
    return utci_all, mrt_all, cat_all


def summarise_utci(utci_all, datetimes):
    """Print and return descriptive statistics."""
    T, N = utci_all.shape
    hours = np.array([d.hour if hasattr(d, 'hour') else pd.Timestamp(d).hour for d in datetimes])

    campus_mean = utci_all.mean(axis=1)
    station_mean = utci_all.mean(axis=0)

    # Stress category distribution
    cat_counts = {}
    for lo, hi, label in UTCI_CATS:
        mask = (utci_all >= lo) & (utci_all < hi)
        pct = mask.sum() / utci_all.size * 100
        cat_counts[label] = round(pct, 1)

    # Diurnal profile
    diurnal = np.zeros(24)
    for h in range(24):
        diurnal[h] = campus_mean[hours == h].mean()

    # Inter-station variability
    hourly_range = utci_all.max(axis=1) - utci_all.min(axis=1)

    stats = {
        'overall_mean': round(float(utci_all.mean()), 2),
        'overall_std': round(float(utci_all.std()), 2),
        'overall_min': round(float(utci_all.min()), 2),
        'overall_max': round(float(utci_all.max()), 2),
        'stress_category_pct': cat_counts,
        'diurnal_profile': [round(float(d), 2) for d in diurnal],
        'station_mean_range': round(float(station_mean.max() - station_mean.min()), 2),
        'hourly_interstation_range_mean': round(float(hourly_range.mean()), 2),
        'hourly_interstation_range_max': round(float(hourly_range.max()), 2),
        'peak_hour': int(np.argmax(diurnal)),
        'trough_hour': int(np.argmin(diurnal)),
    }

    print(f"\n  UTCI Summary:")
    print(f"    Mean: {stats['overall_mean']:.1f} ± {stats['overall_std']:.1f} °C")
    print(f"    Range: {stats['overall_min']:.1f} to {stats['overall_max']:.1f} °C")
    print(f"    Peak at {stats['peak_hour']:02d}:00 ({diurnal[stats['peak_hour']]:.1f} °C)")
    print(f"    Trough at {stats['trough_hour']:02d}:00 ({diurnal[stats['trough_hour']]:.1f} °C)")
    print(f"    Inter-station range: mean {stats['hourly_interstation_range_mean']:.1f} °C, max {stats['hourly_interstation_range_max']:.1f} °C")
    print(f"    Stress categories: {cat_counts}")

    return stats


def eval6_utci_spatial_interpolation(data, coords, datetimes, embeddings, utci_all):
    """
    Hold out same 5 stations as Eval 1. Interpolate their UTCI from
    neighbour embeddings via linear probe. This is the key result:
    can the 6-d embedding predict a derived thermal comfort index
    that was never part of training?
    """
    print("\n" + "=" * 60)
    print("EVAL 6: UTCI Spatial Interpolation")
    print("=" * 60)

    T, N, V = data.shape
    t_tr = int(T * 0.7)
    t_te = int(T * 0.85)

    holdout_idx = [4, 12, 20, 30, 37]  # WS05, WS13, WS21, WS31, WS38 — same as Eval 1
    train_idx = [i for i in range(N) if i not in holdout_idx]

    print(f"  Hold-out: {['WS{:02d}'.format(i+1) for i in holdout_idx]}")
    print(f"  Training: {len(train_idx)} stations")

    # ── Method A: k-NN embedding average + linear probe ──────────────────────
    # Train a Ridge regression: embedding → UTCI
    # Use training stations, training time period
    train_emb = embeddings[:t_tr, train_idx, :].reshape(-1, embeddings.shape[-1])
    train_utci = utci_all[:t_tr, train_idx].reshape(-1)

    probe = Ridge(alpha=1.0)
    probe.fit(train_emb, train_utci)

    # For held-out stations, use k-NN averaged embeddings
    train_coords = coords[train_idx]
    holdout_coords = coords[holdout_idx]
    nn_model = NearestNeighbors(n_neighbors=5).fit(train_coords)
    _, nn_idx = nn_model.kneighbors(holdout_coords)

    test_embeddings = embeddings[t_te:]  # (T_test, N, d)
    test_utci = utci_all[t_te:]          # (T_test, N)

    results = {}
    all_pred = []
    all_true = []

    for hi, ho_station in enumerate(holdout_idx):
        # Average embeddings of 5 nearest training neighbours
        neighbour_stations = [train_idx[j] for j in nn_idx[hi]]
        neigh_emb = test_embeddings[:, neighbour_stations, :].mean(axis=1)  # (T_test, d)

        # Predict UTCI via linear probe
        pred_utci = probe.predict(neigh_emb)
        true_utci = test_utci[:, ho_station]

        mae = float(mean_absolute_error(true_utci, pred_utci))
        rmse = float(np.sqrt(mean_squared_error(true_utci, pred_utci)))
        r2 = float(r2_score(true_utci, pred_utci))

        station_name = f'WS{ho_station+1:02d}'
        results[station_name] = {'MAE': mae, 'RMSE': rmse, 'R2': r2}
        print(f"    {station_name}: MAE={mae:.2f}°C  RMSE={rmse:.2f}°C  R²={r2:.4f}")

        all_pred.append(pred_utci)
        all_true.append(true_utci)

    # Overall
    all_pred = np.concatenate(all_pred)
    all_true = np.concatenate(all_true)
    results['average'] = {
        'MAE': round(float(mean_absolute_error(all_true, all_pred)), 3),
        'RMSE': round(float(np.sqrt(mean_squared_error(all_true, all_pred))), 3),
        'R2': round(float(r2_score(all_true, all_pred)), 4),
    }
    print(f"\n    Average: MAE={results['average']['MAE']:.2f}°C  "
          f"RMSE={results['average']['RMSE']:.2f}°C  R²={results['average']['R2']:.4f}")

    # ── Baseline: k-NN raw variable average → recompute UTCI ─────────────────
    print("\n  Baseline: k-NN raw variable interpolation → UTCI")
    solar_alt, _ = compute_solar_position(datetimes)
    solar_zenith = 90.0 - solar_alt

    baseline_results = {}
    bl_all_pred = []
    bl_all_true = []

    for hi, ho_station in enumerate(holdout_idx):
        neighbour_stations = [train_idx[j] for j in nn_idx[hi]]
        neigh_data = data[t_te:, neighbour_stations, :].mean(axis=1)  # (T_test, 6)

        # Recompute UTCI from averaged raw variables
        ghi = neigh_data[:, 5]
        dni, dhi = decompose_irradiance(ghi, solar_zenith[t_te:], datetimes[t_te:])
        dmrt = compute_delta_mrt(solar_alt[t_te:], dni)
        pred_utci_bl, _, _ = compute_utci_station(
            neigh_data[:, 2], neigh_data[:, 3], neigh_data[:, 0], dmrt
        )
        true_utci = test_utci[:, ho_station]

        mae = float(mean_absolute_error(true_utci, pred_utci_bl))
        rmse = float(np.sqrt(mean_squared_error(true_utci, pred_utci_bl)))
        r2 = float(r2_score(true_utci, pred_utci_bl))

        station_name = f'WS{ho_station+1:02d}'
        baseline_results[station_name] = {'MAE': mae, 'RMSE': rmse, 'R2': r2}
        print(f"    {station_name}: MAE={mae:.2f}°C  RMSE={rmse:.2f}°C  R²={r2:.4f}")

        bl_all_pred.append(pred_utci_bl)
        bl_all_true.append(true_utci)

    bl_all_pred = np.concatenate(bl_all_pred)
    bl_all_true = np.concatenate(bl_all_true)
    baseline_results['average'] = {
        'MAE': round(float(mean_absolute_error(bl_all_true, bl_all_pred)), 3),
        'RMSE': round(float(np.sqrt(mean_squared_error(bl_all_true, bl_all_pred))), 3),
        'R2': round(float(r2_score(bl_all_true, bl_all_pred)), 4),
    }
    print(f"\n    Baseline avg: MAE={baseline_results['average']['MAE']:.2f}°C  "
          f"RMSE={baseline_results['average']['RMSE']:.2f}°C  R²={baseline_results['average']['R2']:.4f}")

    return {
        'embedding_probe': results,
        'raw_knn_baseline': baseline_results,
    }


def eval6_utci_linear_probe(embeddings, utci_all):
    """
    Simpler evaluation: can a linear model predict UTCI from the
    6-d embedding? R² here measures how much comfort information
    the latent space preserves.
    """
    print("\n" + "=" * 60)
    print("EVAL 6b: UTCI Linear Probe (all stations)")
    print("=" * 60)

    T, N, d = embeddings.shape
    t_tr = int(T * 0.7)
    t_te = int(T * 0.85)

    tr_emb = embeddings[:t_tr].reshape(-1, d)
    tr_utci = utci_all[:t_tr].reshape(-1)
    te_emb = embeddings[t_te:].reshape(-1, d)
    te_utci = utci_all[t_te:].reshape(-1)

    probe = Ridge(alpha=1.0)
    probe.fit(tr_emb, tr_utci)
    pred = probe.predict(te_emb)

    results = {
        'MAE': round(float(mean_absolute_error(te_utci, pred)), 3),
        'RMSE': round(float(np.sqrt(mean_squared_error(te_utci, pred))), 3),
        'R2': round(float(r2_score(te_utci, pred)), 4),
    }
    print(f"  Linear probe UTCI: MAE={results['MAE']:.2f}°C  "
          f"RMSE={results['RMSE']:.2f}°C  R²={results['R2']:.4f}")

    return results


def make_figures(utci_all, mrt_all, datetimes, coords, embeddings, tc_results):
    """Generate figures 9, 10, 11 for thermal comfort evaluation."""
    import matplotlib
    matplotlib.use('Agg')
    import matplotlib.pyplot as plt
    import seaborn as sns

    os.makedirs(FIG_DIR, exist_ok=True)

    # Palette
    C = ['#264653', '#2a9d8f', '#e9c46a', '#e76f51']
    sns.set_theme(style='whitegrid', font_scale=1.1)

    T, N = utci_all.shape
    hours = np.array([pd.Timestamp(d).hour for d in datetimes])

    # ── Fig 9: Diurnal UTCI profile with inter-station spread ────────────────
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 5))

    # Panel A: diurnal profile
    diurnal_mean = np.zeros(24)
    diurnal_q25 = np.zeros(24)
    diurnal_q75 = np.zeros(24)
    diurnal_min = np.zeros(24)
    diurnal_max = np.zeros(24)
    for h in range(24):
        vals = utci_all[hours == h].flatten()
        diurnal_mean[h] = vals.mean()
        diurnal_q25[h] = np.percentile(vals, 25)
        diurnal_q75[h] = np.percentile(vals, 75)
        diurnal_min[h] = np.percentile(vals, 5)
        diurnal_max[h] = np.percentile(vals, 95)

    hh = np.arange(24)
    ax1.fill_between(hh, diurnal_min, diurnal_max, alpha=0.15, color=C[3], label='5th–95th pct')
    ax1.fill_between(hh, diurnal_q25, diurnal_q75, alpha=0.3, color=C[1], label='25th–75th pct')
    ax1.plot(hh, diurnal_mean, color=C[0], lw=2, label='Campus mean')

    # Stress thresholds
    ax1.axhline(26, ls='--', color='grey', alpha=0.5, lw=0.8)
    ax1.axhline(32, ls='--', color='orange', alpha=0.5, lw=0.8)
    ax1.axhline(38, ls='--', color='red', alpha=0.5, lw=0.8)
    ax1.text(23.5, 26.3, 'moderate', ha='right', fontsize=8, color='grey')
    ax1.text(23.5, 32.3, 'strong', ha='right', fontsize=8, color='orange')
    ax1.text(23.5, 38.3, 'very strong', ha='right', fontsize=8, color='red')

    ax1.set_xlabel('Hour of day')
    ax1.set_ylabel('UTCI (°C)')
    ax1.set_title('(a) Diurnal UTCI profile')
    ax1.legend(loc='upper left', fontsize=9)
    ax1.set_xlim(0, 23)
    ax1.set_xticks([0, 3, 6, 9, 12, 15, 18, 21])

    # Panel B: stress category distribution
    cats_order = ['no thermal stress', 'moderate heat stress',
                  'strong heat stress', 'very strong heat stress',
                  'extreme heat stress']
    cat_colors = ['#2a9d8f', '#e9c46a', '#e76f51', '#c1121f', '#780000']
    cat_counts = []
    for cat in cats_order:
        for lo, hi, label in UTCI_CATS:
            if label == cat:
                pct = ((utci_all >= lo) & (utci_all < hi)).sum() / utci_all.size * 100
                cat_counts.append(pct)
                break

    bars = ax2.barh(range(len(cats_order)), cat_counts, color=cat_colors[:len(cats_order)])
    ax2.set_yticks(range(len(cats_order)))
    ax2.set_yticklabels([c.capitalize() for c in cats_order], fontsize=9)
    ax2.set_xlabel('Frequency (%)')
    ax2.set_title('(b) UTCI stress categories')
    for i, v in enumerate(cat_counts):
        if v > 1:
            ax2.text(v + 0.5, i, f'{v:.1f}%', va='center', fontsize=9)

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

    # ── Fig 10: Spatial UTCI map (station mean) ──────────────────────────────
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 5))

    station_utci_mean = utci_all.mean(axis=0)
    station_utci_day = utci_all[(hours >= 10) & (hours <= 16)].mean(axis=0) if ((hours >= 10) & (hours <= 16)).sum() > 0 else station_utci_mean

    sc1 = ax1.scatter(coords[:, 1], coords[:, 0], c=station_utci_mean,
                      cmap='RdYlBu_r', s=80, edgecolor='k', linewidth=0.5)
    plt.colorbar(sc1, ax=ax1, label='Mean UTCI (°C)')
    ax1.set_xlabel('Longitude')
    ax1.set_ylabel('Latitude')
    ax1.set_title('(a) Annual mean UTCI')
    ax1.ticklabel_format(useOffset=False)

    sc2 = ax2.scatter(coords[:, 1], coords[:, 0], c=station_utci_day,
                      cmap='RdYlBu_r', s=80, edgecolor='k', linewidth=0.5)
    plt.colorbar(sc2, ax=ax2, label='Mean daytime UTCI (°C)')
    ax2.set_xlabel('Longitude')
    ax2.set_ylabel('Latitude')
    ax2.set_title('(b) Daytime (10:00–16:00) mean UTCI')
    ax2.ticklabel_format(useOffset=False)

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

    # ── Fig 11: UTCI interpolation — predicted vs observed scatter ───────────
    fig, axes = plt.subplots(1, 2, figsize=(12, 5))

    # Get spatial interpolation results for scatter data
    holdout_idx = [4, 12, 20, 30, 37]
    train_idx = [i for i in range(N) if i not in holdout_idx]

    T_data = embeddings.shape[0]
    t_tr = int(T_data * 0.7)
    t_te = int(T_data * 0.85)

    # Embedding probe predictions (re-run for scatter)
    train_coords = coords[train_idx]
    holdout_coords = coords[holdout_idx]
    nn_model = NearestNeighbors(n_neighbors=5).fit(train_coords)
    _, nn_idx = nn_model.kneighbors(holdout_coords)

    train_emb = embeddings[:t_tr, train_idx, :].reshape(-1, embeddings.shape[-1])
    train_utci_flat = utci_all[:t_tr, train_idx].reshape(-1)
    probe = Ridge(alpha=1.0)
    probe.fit(train_emb, train_utci_flat)

    test_utci = utci_all[t_te:]

    scatter_pred_emb = []
    scatter_true_emb = []
    scatter_pred_bl = []
    scatter_true_bl = []

    solar_alt, _ = compute_solar_position(datetimes)
    solar_zenith = 90.0 - solar_alt

    for hi, ho_station in enumerate(holdout_idx):
        neighbour_stations = [train_idx[j] for j in nn_idx[hi]]
        neigh_emb = embeddings[t_te:, neighbour_stations, :].mean(axis=1)
        pred_emb = probe.predict(neigh_emb)
        true = test_utci[:, ho_station]
        scatter_pred_emb.extend(pred_emb)
        scatter_true_emb.extend(true)

        # Baseline
        neigh_data_raw = np.load(f'{CKPT_DIR}/embeddings.npz', allow_pickle=True)['data']
        neigh_raw = neigh_data_raw[t_te:, neighbour_stations, :].mean(axis=1)
        ghi = neigh_raw[:, 5]
        dni, dhi = decompose_irradiance(ghi, solar_zenith[t_te:], datetimes[t_te:])
        dmrt = compute_delta_mrt(solar_alt[t_te:], dni)
        pred_bl, _, _ = compute_utci_station(neigh_raw[:, 2], neigh_raw[:, 3], neigh_raw[:, 0], dmrt)
        scatter_pred_bl.extend(pred_bl)
        scatter_true_bl.extend(true)

    scatter_pred_emb = np.array(scatter_pred_emb)
    scatter_true_emb = np.array(scatter_true_emb)
    scatter_pred_bl = np.array(scatter_pred_bl)
    scatter_true_bl = np.array(scatter_true_bl)

    # Subsample for plotting
    n_plot = min(5000, len(scatter_pred_emb))
    rng = np.random.RandomState(42)
    idx = rng.choice(len(scatter_pred_emb), n_plot, replace=False)

    # Panel A: embedding probe
    ax = axes[0]
    ax.scatter(scatter_true_emb[idx], scatter_pred_emb[idx], alpha=0.15, s=8, c=C[1])
    lims = [min(scatter_true_emb.min(), scatter_pred_emb.min()) - 1,
            max(scatter_true_emb.max(), scatter_pred_emb.max()) + 1]
    ax.plot(lims, lims, 'k--', lw=1, alpha=0.5)
    r2_emb = r2_score(scatter_true_emb, scatter_pred_emb)
    mae_emb = mean_absolute_error(scatter_true_emb, scatter_pred_emb)
    ax.text(0.05, 0.92, f'R² = {r2_emb:.3f}\nMAE = {mae_emb:.2f}°C',
            transform=ax.transAxes, fontsize=10, va='top',
            bbox=dict(boxstyle='round', facecolor='white', alpha=0.8))
    ax.set_xlabel('Observed UTCI (°C)')
    ax.set_ylabel('Predicted UTCI (°C)')
    ax.set_title('(a) Embedding + linear probe')
    ax.set_xlim(lims); ax.set_ylim(lims)
    ax.set_aspect('equal')

    # Panel B: raw kNN baseline
    ax = axes[1]
    ax.scatter(scatter_true_bl[idx], scatter_pred_bl[idx], alpha=0.15, s=8, c=C[3])
    r2_bl = r2_score(scatter_true_bl, scatter_pred_bl)
    mae_bl = mean_absolute_error(scatter_true_bl, scatter_pred_bl)
    ax.plot(lims, lims, 'k--', lw=1, alpha=0.5)
    ax.text(0.05, 0.92, f'R² = {r2_bl:.3f}\nMAE = {mae_bl:.2f}°C',
            transform=ax.transAxes, fontsize=10, va='top',
            bbox=dict(boxstyle='round', facecolor='white', alpha=0.8))
    ax.set_xlabel('Observed UTCI (°C)')
    ax.set_ylabel('Predicted UTCI (°C)')
    ax.set_title('(b) Raw k-NN interpolation')
    ax.set_xlim(lims); ax.set_ylim(lims)
    ax.set_aspect('equal')

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


def run_all():
    """Run the complete thermal comfort evaluation."""
    print("=" * 60)
    print("EVALUATION 6: Thermal Comfort (UTCI)")
    print("=" * 60)

    # Load data and pre-trained embeddings
    data, coords, datetimes = load_nus40(DATA_DIR)
    npz = np.load(f'{CKPT_DIR}/embeddings.npz', allow_pickle=True)
    embeddings = npz['embeddings']

    datetimes_list = pd.to_datetime(datetimes)

    # ── Step 1: Compute UTCI at all 40 stations ─────────────────────────────
    print("\nStep 1: Computing UTCI for all 40 stations...")
    utci_all, mrt_all, cat_all = compute_all_stations(data, datetimes_list)

    # ── Step 2: Summarise ────────────────────────────────────────────────────
    print("\nStep 2: Summary statistics")
    stats = summarise_utci(utci_all, datetimes_list)

    # ── Step 3: Spatial interpolation evaluation ─────────────────────────────
    print("\nStep 3: Spatial interpolation of UTCI")
    interp_results = eval6_utci_spatial_interpolation(
        data, coords, datetimes_list, embeddings, utci_all)

    # ── Step 4: Linear probe evaluation ──────────────────────────────────────
    print("\nStep 4: Linear probe")
    probe_results = eval6_utci_linear_probe(embeddings, utci_all)

    # ── Step 5: Save results ─────────────────────────────────────────────────
    tc_results = {
        'summary': stats,
        'spatial_interpolation': interp_results,
        'linear_probe': probe_results,
        'parameters': {
            'z_measurement': Z_MEAS,
            'z0_roughness': Z0,
            'sky_view_factor': F_SVV,
            'body_exposure': F_BES,
            'campus_lat': CAMPUS_LAT,
            'campus_lon': CAMPUS_LON,
        }
    }

    os.makedirs(RESULTS_DIR, exist_ok=True)
    with open(f'{RESULTS_DIR}/thermal_comfort.json', 'w') as f:
        json.dump(tc_results, f, indent=2)
    print(f"\n  Results saved to {RESULTS_DIR}/thermal_comfort.json")

    # Save UTCI array
    np.savez_compressed(f'{RESULTS_DIR}/utci_all.npz',
                        utci=utci_all, mrt=mrt_all,
                        datetimes=np.array(datetimes_list.astype(str)))
    print(f"  UTCI array saved to {RESULTS_DIR}/utci_all.npz")

    # ── Step 6: Figures ──────────────────────────────────────────────────────
    print("\nStep 6: Generating figures...")
    make_figures(utci_all, mrt_all, datetimes_list, coords, embeddings, tc_results)

    # ── Update all_results.json ──────────────────────────────────────────────
    all_results_path = f'{RESULTS_DIR}/all_results.json'
    if os.path.exists(all_results_path):
        with open(all_results_path) as f:
            all_results = json.load(f)
    else:
        all_results = {}

    all_results['thermal_comfort'] = tc_results
    with open(all_results_path, 'w') as f:
        json.dump(all_results, f, indent=2)
    print(f"  Updated {all_results_path}")

    print("\n" + "=" * 60)
    print("DONE: Evaluation 6 complete")
    print("=" * 60)

    return tc_results


if __name__ == '__main__':
    run_all()