"""
12 — Rth response surface in PC3-PC4 latent space.

Fits PCA on the first TRAIN_SIZE samples, projects all 2000 onto the
basis, computes per-chip thermal resistance from the field data,
interpolates Rth1 / Rth2 onto a regular PC3-PC4 grid, and plots
side-by-side contour response surfaces with mean-gradient arrows.
"""

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.colors as mcolors
import matplotlib.ticker as mticker
from scipy.interpolate import griddata
from sklearn.decomposition import PCA

from _utils import (
    load_data, compute_rth, TRAIN_SIZE, N_MODES, PLOTS_DIR,
)

GRID_RES = 80
FS = 20
RTH_MIN, RTH_MAX = 0.1, 0.13


def make_norm():
    return mcolors.Normalize(vmin=RTH_MIN, vmax=RTH_MAX)


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

    pca = PCA(n_components=N_MODES).fit(fields[:TRAIN_SIZE])
    coeffs = pca.transform(fields)
    var = pca.explained_variance_ratio_

    rth1, rth2, _, _ = compute_rth(fields, coords, parameters)
    print(f"Rth1: [{rth1.min():.3f}, {rth1.max():.3f}] K/W")
    print(f"Rth2: [{rth2.min():.3f}, {rth2.max():.3f}] K/W")

    pc3 = coeffs[:, 2]; pc4 = coeffs[:, 3]
    pc3v = np.linspace(pc3.min(), pc3.max(), GRID_RES)
    pc4v = np.linspace(pc4.min(), pc4.max(), GRID_RES)
    P3, P4 = np.meshgrid(pc3v, pc4v)
    pts = np.column_stack([pc3, pc4])
    R1g = griddata(pts, rth1, (P3, P4), method="cubic")
    R2g = griddata(pts, rth2, (P3, P4), method="cubic")

    norm = make_norm()
    levels = np.linspace(RTH_MIN, RTH_MAX, 25)
    fig, axes = plt.subplots(1, 2, figsize=(16, 7))

    dpc3 = P3[0, 1] - P3[0, 0]
    dpc4 = P4[1, 0] - P4[0, 0]

    cbars = []
    for ax, Z, lbl in zip(axes, [R1g, R2g],
                           [r"$R_{\mathrm{th},1}$ (K/W)",
                            r"$R_{\mathrm{th},2}$ (K/W)"]):
        ax.contourf(P3, P4, Z, levels=[-1e10, RTH_MIN], colors=["red"], alpha=0.6)
        ax.contourf(P3, P4, Z, levels=[RTH_MAX, 1e10], colors=["red"], alpha=0.6)
        cf = ax.contourf(P3, P4, Z, levels=levels, cmap="viridis",
                         norm=norm, extend="both")
        ax.scatter(pc3, pc4, s=3, alpha=0.15, color="white", edgecolors="none")

        dZ_dpc4, dZ_dpc3 = np.gradient(Z, dpc4, dpc3)
        valid = ~np.isnan(dZ_dpc3) & ~np.isnan(dZ_dpc4)
        gx = np.nanmean(dZ_dpc3[valid])
        gy = np.nanmean(dZ_dpc4[valid])
        dx, dy = -gx, -gy
        mag = np.hypot(dx, dy)
        L = 0.2 * min(P3.max() - P3.min(), P4.max() - P4.min())
        dx, dy = dx / mag * L, dy / mag * L
        cx = (P3.min() + P3.max()) / 2
        cy = (P4.min() + P4.max()) / 2
        ax.annotate("", xy=(cx + dx / 2, cy + dy / 2),
                    xytext=(cx - dx / 2, cy - dy / 2),
                    arrowprops=dict(arrowstyle="->", color="white",
                                    lw=3, mutation_scale=20))

        ax.set_xlabel(f"PC3 ({var[2]*100:.1f}%)", fontsize=FS)
        ax.set_ylabel(f"PC4 ({var[3]*100:.1f}%)", fontsize=FS)
        ax.tick_params(labelsize=FS - 4)
        ax.xaxis.set_major_formatter(mticker.FormatStrFormatter("%g"))
        ax.yaxis.set_major_formatter(mticker.FormatStrFormatter("%g"))
        cb = plt.colorbar(cf, ax=ax, shrink=0.85, extend="both")
        cb.set_label(lbl, fontsize=FS - 2)
        cb.ax.tick_params(labelsize=FS - 4)
        cbars.append(cb)

    plt.tight_layout()
    ticks = np.linspace(RTH_MIN, RTH_MAX, 5)
    for cb in cbars:
        cb.set_ticks(ticks)
        cb.set_ticklabels([f"{t:.3f}" for t in ticks])

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


if __name__ == "__main__":
    main()