"""Dump decoded free-running latent trajectories for the parallel-cot organism.

For a few problems, free-run and decode each latent `z_t` via the per-cell head (argmax per cell),
recording the ground-truth CA row vs the decoded row at every step. Shows the organism tracing the CA
step by step in its latents — including a chain LONGER than any trained. Writes
data/single_traj_<tag>.json (read by the model card) and prints a readable dump.

    python -m latent_threads.eval_single_traj --ckpt .../sg1/diffuse/best
"""
from __future__ import annotations

import argparse
import json
import os

import dotenv

dotenv.load_dotenv()
_USER = os.environ.get("USER", "jbauer")
os.environ.setdefault("HF_HOME", f"/workspace-vast/{_USER}/hf")

import random

import torch

from latent_threads import tasks as LT
from latent_threads.eval_single_report import load
from latent_threads.single import single_forward

REPO_DATA = "/workspace-vast/jbauer/activation_oracles_dev/data"


@torch.no_grad()
def run(ckpt, model_name, device):
    model, head, codebook, q_emb, tok, did, K, cfg, tag = load(ckpt, model_name, device)
    im = tok.convert_tokens_to_ids("<|im_end|>")
    base_m = cfg["task_kwargs"]["m"]
    rng = random.Random(0)
    specs = [("trained", base_m, 2), ("longer-than-trained", base_m + 4, 1)]   # (label, chain length, count)
    out = []
    for label, m, count in specs:
        task = LT.make_task(cfg["task"], k=K, m=m)
        probs = [task.sample(rng) for _ in range(count)]
        _, _, (aux, _) = single_forward(model, head, codebook, q_emb, tok, task, probs, im, device, with_answer=False)
        dec = aux.argmax(-1).tolist()                                          # [B, m, K] head-decoded rows
        for j, p in enumerate(probs):
            steps = [{"gt": list(p.rows[t]), "dec": dec[j][t]} for t in range(m)]
            out.append({"label": label, "m": m, "init": list(p.init), "q": p.q, "steps": steps,
                        "answer_gt": p.rows[-1][p.q], "answer_model": dec[j][m - 1][p.q]})
    os.makedirs(REPO_DATA, exist_ok=True)
    json.dump(out, open(f"{REPO_DATA}/single_traj_{tag}.json", "w"), indent=2)
    for e in out:
        print(f"\n[{e['label']} T={e['m']}] init c(0)={e['init']}  query=c{e['q']+1}", flush=True)
        for t, s in enumerate(e["steps"]):
            print(f"  z_{t+1} -> {s['dec']}   (CA truth {s['gt']})   {'OK' if s['dec']==s['gt'] else 'XX'}", flush=True)
        print(f"  read c{e['q']+1} of z_{e['m']} -> {e['answer_model']}   (CA truth {e['answer_gt']})", flush=True)
    print(f"\nwrote {len(out)} trajectories -> {REPO_DATA}/single_traj_{tag}.json", flush=True)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--ckpt", required=True); ap.add_argument("--model", default="Qwen/Qwen3-8B")
    args = ap.parse_args()
    run(args.ckpt, args.model, "cuda")


if __name__ == "__main__":
    main()