"""
05 — Train/Val MSE vs epoch from the canonical model's training history.
"""

import matplotlib.pyplot as plt

from _utils import load_artifacts, PLOTS_DIR


def main():
    trained = load_artifacts()
    h = trained["history"]
    epochs = range(1, len(h["loss"]) + 1)

    FS = 20
    fig, ax = plt.subplots(figsize=(8, 6))
    ax.plot(epochs, h["loss"], label="Train", linewidth=2)
    if "val_loss" in h:
        ax.plot(epochs, h["val_loss"], label="Val", linewidth=2)
    ax.set_xlabel("Epoch", fontsize=FS)
    ax.set_ylabel("Loss (MSE)", fontsize=FS)
    ax.tick_params(labelsize=FS - 4)
    ax.legend(fontsize=FS - 4)
    ax.grid(True, alpha=0.3)

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


if __name__ == "__main__":
    main()