"""Append one progress row for a mixed-corpus checkpoint: held-out dialogue / held-out varied / old
val2048 (normal + field-alone), train stats, 3 probe chats (unseeded) + seed 1.  1 thread.
usage: tsv_mix.py TSV CKPT RUN_METRICS_JSON [--no-chats]"""
import json, os, re, sys, time
from pathlib import Path
import torch
sys.path.insert(0, str(Path(__file__).resolve().parent))
import mix_common as M


def main():
    torch.set_num_threads(1)
    tsv, ck, rm = sys.argv[1:4]
    t0 = time.time(); m = M.C.load_model(ck); v = M.all_vals(m)
    tr = json.load(open(rm)).get("train", {}) if os.path.exists(rm) else {}
    lam = tr.get("field_leak_learnable_rates")
    chats = [] if "--no-chats" in sys.argv else M.C.probe_chats(m, seeds=(1,))
    step = re.search(r"step_(\d+)_", ck).group(1)
    cols = [("step", step)] + [(k, f"{x:.4f}") for k, x in v.items()] + [
        ("train_final_loss", f"{tr.get('final_loss', float('nan')):.4f}"),
        ("steps_per_s", f"{tr.get('transitions_per_second', 0):.1f}"),
        ("cap_frac", f"{tr.get('field_limited_fraction', float('nan')):.3f}"),
        ("nan", str(tr.get("nan_or_inf"))),
        ("leak_mean", f"{sum(lam) / len(lam):.4f}" if lam else ""),
        ("leak_min", f"{min(lam):.4f}" if lam else ""),
        ("time", time.strftime("%Y-%m-%dT%H:%M:%S%z"))]
    cols += [(f"chat_{r['seed'] or 'u'}_{i % 3}", " ".join(r["resp"].split())[:90].replace("\t", " ")) for i, r in enumerate(chats)]
    new = not os.path.exists(tsv)
    with open(tsv, "a", encoding="utf-8") as fh:
        if new: fh.write("\t".join(c for c, _ in cols) + "\n")
        fh.write("\t".join(x for _, x in cols) + "\n")
    print("\t".join(x for _, x in cols), f"({time.time() - t0:.0f}s)", flush=True)


if __name__ == "__main__":
    main()