import csv
import matplotlib.pyplot as plt

epochs, tr, va = [], [], []
with open("checkpoints_diag/history.csv") as f:
    for row in csv.DictReader(f):
        epochs.append(int(row["epoch"]))
        tr.append(float(row["train_mse"]))
        va.append(float(row["val_mse"]))

plt.figure(figsize=(7, 4))
plt.plot(epochs, tr, label="train MSE", marker="o")
plt.plot(epochs, va, label="val MSE", marker="s")
plt.xlabel("epoch"); plt.ylabel("MSE"); plt.legend(); plt.grid(alpha=0.3)
plt.title("Training vs. validation loss")
plt.tight_layout()
plt.savefig("outputs/loss_curves.png", dpi=120)
print("saved outputs/loss_curves.png")