"""
03 — Original vs 1-PC vs 4-PC reconstruction (sample 329).

Quad plot: original field, 1-PC reconstruction, 4-PC reconstruction,
shared horizontal colorbar.
"""

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.cm as mcm
import matplotlib.colors as mcolors
from sklearn.decomposition import PCA

from _utils import load_data, PLOTS_DIR

SAMPLE_INDEX = 329
N_MODES_FIT = 4


def reconstruct(pca, field, n_keep):
    scores = pca.transform(field[np.newaxis, :])
    trunc = np.zeros_like(scores)
    trunc[:, :n_keep] = scores[:, :n_keep]
    return pca.inverse_transform(trunc)[0]


def main():
    _, fields, coords, _ = load_data()
    pca = PCA(n_components=N_MODES_FIT).fit(fields)

    y, z = coords[:, 1], coords[:, 2]
    original = fields[SAMPLE_INDEX]
    recon1 = reconstruct(pca, original, 1)
    recon4 = reconstruct(pca, original, 4)

    all_vals = np.concatenate([original, recon1, recon4])
    vmin, vmax = float(all_vals.min()), float(all_vals.max())
    norm = mcolors.Normalize(vmin=vmin, vmax=vmax)
    sm = mcm.ScalarMappable(cmap="inferno", norm=norm)
    sm.set_array([])
    kw = dict(cmap="inferno", s=1.5, vmin=vmin, vmax=vmax, rasterized=True)

    fig, axes = plt.subplots(2, 2, figsize=(8, 4))
    panels = [
        (axes[0, 0], original, "Original"),
        (axes[0, 1], recon1,   "1-PC Reconstruction"),
        (axes[1, 0], recon4,   "4-PC Reconstruction"),
    ]
    for ax, field, title in panels:
        ax.scatter(z, y, c=field, **kw)
        ax.invert_xaxis()
        ax.set_xticks([]); ax.set_yticks([])
        ax.set_aspect("equal")
        ax.set_title(title, fontsize=18)

    ax_cb = axes[1, 1]
    ax_cb.axis("off")
    pos = ax_cb.get_position()
    cbar_ax = fig.add_axes((pos.x0 + pos.width * 0.05,
                            pos.y0 + pos.height * 0.38,
                            pos.width * 0.9,
                            pos.height * 0.25))
    cbar = fig.colorbar(sm, cax=cbar_ax, orientation="horizontal")
    ticks = [round(v) for v in np.linspace(vmin, vmax, 5)]
    cbar.set_ticks(ticks); cbar.set_ticklabels([str(t) for t in ticks])
    cbar_ax.tick_params(labelsize=12)
    cbar_ax.set_title("Temperature (K)", fontsize=18, pad=8)

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


if __name__ == "__main__":
    main()