"""
07 — Single-sample prediction comparison plots.

Loads the canonical model, picks 5 evenly spaced indices from the
held-out test set, generates one figure per sample with three panels
(prediction / truth / abs-error) and the input parameters listed.
"""

import time
import numpy as np
import matplotlib.pyplot as plt
from sklearn.metrics import r2_score, mean_absolute_error

from _utils import (
    load_data, load_artifacts, predict_field, yz_mm,
    TEST_SIZE, PLOTS_DIR,
)

N_PLOTS = 5
DISPLAY = {
    "chip1.thermal.heat_flux":        ("Chip 1 Heat Flux",     "W/m²", ".0f"),
    "chip2.thermal.heat_flux":        ("Chip 2 Heat Flux",     "W/m²", ".0f"),
    "inlet1.momentum.mass_flow_rate": ("Inlet 1 Mass Flow",    "kg/s",      ".6f"),
    "inlet1.thermal.total_temperature": ("Inlet 1 Temperature", "K",         ".1f"),
    "inlet2.momentum.mass_flow_rate": ("Inlet 2 Mass Flow",    "kg/s",      ".6f"),
    "inlet2.thermal.total_temperature": ("Inlet 2 Temperature", "K",         ".1f"),
}


def plot_one(coords, y_pred, y_true, idx, t_inf_ms, x_vals, p_names, out_dir):
    FS = 16
    y_mm, z_mm = yz_mm(coords)
    abs_err = np.abs(y_pred - y_true)
    vmin = min(y_true.min(), y_pred.min())
    vmax = max(y_true.max(), y_pred.max())

    fig = plt.figure(figsize=(14, 12))
    gs = fig.add_gridspec(3, 2, width_ratios=[3, 1], wspace=0.05, hspace=0.15)
    ax_p, ax_t, ax_e = (fig.add_subplot(gs[i, 0]) for i in range(3))
    ax_text = fig.add_subplot(gs[:, 1])

    for ax, vals, title, cmap, lbl, kw in [
        (ax_p, y_pred,  "Prediction",     "jet", "T (K)",        dict(vmin=vmin, vmax=vmax)),
        (ax_t, y_true,  "Truth",          "jet", "T (K)",        dict(vmin=vmin, vmax=vmax)),
        (ax_e, abs_err, "Absolute Error", "hot", "|Error| (K)",  dict()),
    ]:
        sc = ax.scatter(z_mm, y_mm, c=vals, s=1, cmap=cmap, **kw)
        ax.set_xlabel("Z (mm)", fontsize=FS)
        ax.set_ylabel("Y (mm)", fontsize=FS)
        ax.set_title(title, fontsize=FS)
        ax.tick_params(labelsize=FS - 4)
        ax.set_aspect("equal")
        cb = plt.colorbar(sc, ax=ax, orientation="vertical", shrink=0.8)
        cb.set_label(lbl, fontsize=FS - 2)
        cb.ax.tick_params(labelsize=FS - 4)

    ax_text.axis("off")
    r2  = r2_score(y_true, y_pred)
    mae = mean_absolute_error(y_true, y_pred)

    lines = [f"Test Sample {idx}", "", "Input Parameters", "-" * 28]
    for name, v in zip(p_names, x_vals):
        if name in DISPLAY:
            lbl, unit, fmt = DISPLAY[name]
            lines.append(f"{lbl}:")
            lines.append(f"  {v:{fmt}} {unit}")
        else:
            lines.append(f"{name}: {v:.6g}")
    lines += ["", "Metrics", "-" * 28,
              f"R²:        {r2:.6f}",
              f"MAE:       {mae:.3f} K",
              f"Max AE:    {abs_err.max():.3f} K",
              f"Inference: {t_inf_ms:.2f} ms"]
    ax_text.text(0.05, 0.95, "\n".join(lines),
                 transform=ax_text.transAxes, fontsize=FS - 4,
                 verticalalignment="top", fontfamily="monospace")

    plt.tight_layout()
    out = out_dir / f"comparison_sample_{idx:04d}.png"
    plt.savefig(out, dpi=300, bbox_inches="tight")
    print(f"  Saved {out.name}")
    plt.close()


def main():
    parameters, fields, coords, p_names = load_data()
    X_test = parameters[-TEST_SIZE:]
    Y_test = fields[-TEST_SIZE:]

    trained = load_artifacts()

    out_dir = PLOTS_DIR / "single_predictions"
    out_dir.mkdir(parents=True, exist_ok=True)

    indices = np.linspace(0, TEST_SIZE - 1, N_PLOTS, dtype=int)
    for idx in indices:
        x = X_test[idx:idx+1]
        t0 = time.perf_counter()
        y_pred = predict_field(trained, x)[0]
        t_inf_ms = (time.perf_counter() - t0) * 1000
        plot_one(coords, y_pred, Y_test[idx], int(idx),
                 t_inf_ms, X_test[idx], p_names, out_dir)


if __name__ == "__main__":
    main()