"""
02 — Plot the first 6 POD spatial mode shapes on the yz-mid plane.

Refits PCA with 6 components on the full dataset (the canonical model
in 01 only retains N_MODES=4, but for visualization we want a couple
extra modes to illustrate the diminishing structure).
"""

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

from _utils import load_data, PLOTS_DIR

N_PLOT = 6


def main():
    _, fields, coords, _ = load_data()

    pca = PCA(n_components=N_PLOT).fit(fields)
    var = pca.explained_variance_ratio_

    y, z = coords[:, 1], coords[:, 2]

    fig, axes = plt.subplots(3, 2, figsize=(8, 7))
    for i, ax in enumerate(axes.flatten()):
        mode = pca.components_[i]
        vmax = np.abs(mode).max()
        ax.scatter(z, y, c=mode, cmap="RdBu_r", s=1.5,
                   vmin=-vmax, vmax=vmax, rasterized=True)
        ax.invert_xaxis()
        ax.set_xticks([]); ax.set_yticks([])
        ax.set_aspect("equal")
        ax.set_title(f"PC{i+1}  |  Var: {var[i]*100:.2f}%", fontsize=14)

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


if __name__ == "__main__":
    main()