"""
08 — Mean absolute error per spatial point across the test set.

Loads the canonical model, runs inference on all 400 held-out test
samples, computes |error| at every spatial point per sample, averages
across the test set, plots the resulting MAE field on the yz-mid plane.
"""

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,
)


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

    trained = load_artifacts()
    Y_pred = predict_field(trained, X_test)

    abs_err = np.abs(Y_pred - Y_test)        # (400, n_pts)
    mean_err = abs_err.mean(axis=0)          # (n_pts,)

    r2_all  = r2_score(Y_test.flatten(), Y_pred.flatten())
    mae_all = mean_absolute_error(Y_test.flatten(), Y_pred.flatten())
    print(f"Test set: R²={r2_all:.6f}  MAE={mae_all:.4f} K  "
          f"MaxAE={abs_err.max():.4f} K")
    print(f"MAE field: mean={mean_err.mean():.4f} K  "
          f"max={mean_err.max():.4f} K  min={mean_err.min():.4f} K")

    y_mm, z_mm = yz_mm(coords)

    FS = 20
    fig, ax = plt.subplots(figsize=(10, 6))
    sc = ax.scatter(z_mm, y_mm, c=mean_err, s=1, cmap="hot")
    ax.set_xlabel("Z (mm)", fontsize=FS)
    ax.set_ylabel("Y (mm)", fontsize=FS)
    ax.set_aspect("equal")
    ax.tick_params(labelsize=FS - 4)
    cb = plt.colorbar(sc, ax=ax, orientation="vertical", shrink=0.5)
    cb.set_label("Mean |Error| (K)", fontsize=FS - 2)
    cb.ax.tick_params(labelsize=FS - 4)

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


if __name__ == "__main__":
    main()