"""Event file -> loss_curves.png (diff/repa/lr + fid). watcher가 매 사이클 갱신."""
import glob, os
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from tensorboard.backend.event_processing.event_accumulator import EventAccumulator

BASE = os.path.join(os.path.dirname(os.path.abspath(__file__)),
                    "output/tokenizer/models_l_spatial_phase2")
LOGDIR = os.path.join(BASE, "logs/semanticist")
OUT = os.path.join(BASE, "loss_curves.png")


def load():
    ea = EventAccumulator(sorted(glob.glob(os.path.join(LOGDIR, "events*")))[-1])
    ea.Reload()
    d = {}
    for t in ea.Tags()["scalars"]:
        s = ea.Scalars(t)
        d[t] = ([x.step for x in s], [x.value for x in s])
    return d


def main():
    d = load()
    fig, ax = plt.subplots(2, 2, figsize=(13, 8))
    if "diff_loss" in d:
        ax[0, 0].plot(*d["diff_loss"], lw=0.7, color="C0")
        ax[0, 0].set_title(f"diff_loss (last {d['diff_loss'][1][-1]:.4f})")
        ax[0, 0].set_xlabel("step"); ax[0, 0].grid(alpha=.3)
    if "repa_loss" in d:
        ax[0, 1].plot(*d["repa_loss"], lw=0.7, color="C3")
        ax[0, 1].set_title(f"repa_loss (last {d['repa_loss'][1][-1]:.4f})")
        ax[0, 1].set_xlabel("step"); ax[0, 1].grid(alpha=.3)
    if "lr" in d:
        ax[1, 0].plot(*d["lr"], lw=0.9, color="C2")
        ax[1, 0].set_title("lr (warmup)")
        ax[1, 0].set_xlabel("step"); ax[1, 0].grid(alpha=.3)
    if "fid" in d and len(d["fid"][0]) > 0:
        ax[1, 1].plot(*d["fid"], "o-", color="C1", label="fid")
        ax[1, 1].set_title(f"eval fid (last {d['fid'][1][-1]:.2f} @{d['fid'][0][-1]})")
        ax[1, 1].set_xlabel("step"); ax[1, 1].grid(alpha=.3)
    last_step = d.get("diff_loss", ([0], [0]))[0][-1]
    fig.suptitle(f"SpatialDiffuseSlot-L Phase2 — step {last_step}", fontsize=13)
    fig.tight_layout()
    fig.savefig(OUT, dpi=110)
    print("wrote", OUT, "@step", last_step)


if __name__ == "__main__":
    main()