"""Single-token organism training curve: free-running readout + per-cell state + curriculum growth.

    python -m latent_threads.plot_single_summary <train_log> [tag]
"""
import re
import sys

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

RES = "/workspace-vast/jbauer/exp/latent_threads/results"
log = sys.argv[1]
tag = sys.argv[2] if len(sys.argv) > 2 else "sg1_diffuse_best"
# [eval] step N cur_max=C readout@cur=A readout@6=B state@cur=S ... (tf_p=T, ...)
EV = re.compile(r"\[eval\] step (\d+) cur_max=(\d+) readout@cur=([\d.]+) readout@\d+=([\d.]+) state@cur=([\d.]+).*tf_p=([\d.]+)")

steps, cmax, ro, st, tf = [], [], [], [], []
for line in open(log, errors="replace"):
    m = EV.search(line)
    if m:
        steps.append(int(m[1])); cmax.append(int(m[2])); ro.append(float(m[4])); st.append(float(m[5])); tf.append(float(m[6]))

cmax_max = max(cmax) if cmax else 6
fig, ax = plt.subplots(figsize=(7.5, 4.5))
ax.plot(steps, ro, marker=".", label="free-running readout @ target length")
ax.plot(steps, st, marker=".", label="per-cell state decodability")
ax.plot(steps, tf, ls="--", color="gray", alpha=0.6, label="teacher-forcing prob (anneal)")
ax.plot(steps, [c / cmax_max for c in cmax], ls="-.", color="purple", alpha=0.7, label=f"curriculum length / {cmax_max}")
ax.axhline(0.1, ls=":", color="red", alpha=0.5)
ax.axhline(0.9, ls=":", color="green", alpha=0.5)
ax.set_xlabel("training step"); ax.set_ylabel("accuracy / probability"); ax.set_ylim(0, 1.02)
ax.legend(fontsize=8); ax.set_title("Single-token latent CoT: readout + state vs TF anneal & curriculum")
png = f"{RES}/single_training_{tag}.png"
fig.savefig(png, bbox_inches="tight", dpi=130)
print(png)