"""
10 — POD mode α sweep heatmaps + per-mode delta map.

Fits PCA on the first TRAIN_SIZE samples; for each of the 4 retained
PCs sweeps its coefficient α from min to max in 5 steps with the other
coefficients held at their means, reconstructs the field at each step.

Outputs:
  - pod_alpha_sweep.png — 4×5 sweep grid (paper Fig 8)
  - pod_mode_delta.png  — 2×2 (α=0%) − (α=100%) field per mode (paper Fig 9)
"""

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.colors import Normalize
from sklearn.decomposition import PCA

from _utils import load_data, TRAIN_SIZE, N_MODES, PLOTS_DIR

N_STEPS = 5
FS = 30


def main():
    _, fields, coords, _ = load_data()
    fields = fields[:TRAIN_SIZE]

    pca = PCA(n_components=N_MODES).fit(fields)
    coeffs = pca.transform(fields)
    coeff_min  = coeffs.min(axis=0)
    coeff_max  = coeffs.max(axis=0)
    coeff_mean = coeffs.mean(axis=0)
    var = pca.explained_variance_ratio_

    y, z = coords[:, 1], coords[:, 2]
    norm = Normalize(vmin=fields.min(), vmax=fields.max())
    fractions = np.linspace(0, 1, N_STEPS)

    fig, axes = plt.subplots(N_MODES, N_STEPS, figsize=(4 * N_STEPS, 2 * N_MODES))
    sc = None
    for m in range(N_MODES):
        sweep = coeff_min[m] + fractions * (coeff_max[m] - coeff_min[m])
        for k, val in enumerate(sweep):
            ax = axes[m, k]
            c = coeff_mean.copy()
            c[m] = val
            field_rec = pca.mean_ + c @ pca.components_
            sc = ax.scatter(z, y, c=field_rec, s=0.3, cmap="jet", norm=norm)
            ax.set_aspect("equal")
            ax.set_xticks([]); ax.set_yticks([])
            if m == 0:
                ax.set_title(rf"$\alpha$ = {fractions[k]*100:.0f}%",
                             fontsize=FS - 4)
            if k == 0:
                ax.set_ylabel(f"PC{m+1}\n({var[m]*100:.1f}%)", fontsize=FS - 4)

    fig.subplots_adjust(right=0.92, hspace=0.15, wspace=0.05)
    cbar_ax = fig.add_axes([0.93, 0.15, 0.012, 0.7])
    cb = fig.colorbar(sc, cax=cbar_ax)
    cb.set_label("T (K)", fontsize=FS - 2)
    cb.ax.tick_params(labelsize=FS - 6)

    PLOTS_DIR.mkdir(parents=True, exist_ok=True)
    out = PLOTS_DIR / "pod_alpha_sweep.png"
    plt.savefig(out, dpi=300, bbox_inches="tight")
    print(f"Saved {out}")
    plt.close()

    # --- Per-mode delta plot (paper Fig 9): (α=0%) - (α=100%) per PC ---
    deltas = []
    for m in range(N_MODES):
        c_lo = coeff_mean.copy(); c_lo[m] = coeff_min[m]
        c_hi = coeff_mean.copy(); c_hi[m] = coeff_max[m]
        f_lo = pca.mean_ + c_lo @ pca.components_
        f_hi = pca.mean_ + c_hi @ pca.components_
        deltas.append(f_lo - f_hi)
    vabs = max(np.abs(d).max() for d in deltas)
    delta_norm = Normalize(vmin=-vabs, vmax=vabs)

    fig, axes = plt.subplots(2, 2, figsize=(12, 6))
    axes = axes.flatten()
    sc = None
    for m in range(N_MODES):
        ax = axes[m]
        sc = ax.scatter(z, y, c=deltas[m], s=0.5, cmap="coolwarm",
                        norm=delta_norm)
        ax.set_aspect("equal")
        ax.set_xticks([]); ax.set_yticks([])
        ax.set_title(f"PC{m+1} ({var[m]*100:.1f}%)", fontsize=FS - 10)

    fig.subplots_adjust(right=0.88, hspace=0.15, wspace=0.05)
    cbar_ax = fig.add_axes([0.90, 0.15, 0.02, 0.7])
    cb = fig.colorbar(sc, cax=cbar_ax)
    cb.set_label(r"$\Delta T$ (K)", fontsize=FS - 10)
    cb.ax.tick_params(labelsize=FS - 12)

    out_d = PLOTS_DIR / "pod_mode_delta.png"
    plt.savefig(out_d, dpi=300, bbox_inches="tight")
    print(f"Saved {out_d}")
    plt.close()


if __name__ == "__main__":
    main()