#!/usr/bin/env python3
"""
Observability: dump the full recurrent trajectory for one position.

Emits a human-readable trace and/or machine-readable JSON containing
h0..hn, Δh_t, memory, routing weights, halting probabilities, per-step move
latents and candidate move scores.

    python inspect_trace.py --checkpoint runs/x/latest.pt --fen "..." --steps 8
"""
from __future__ import annotations

import argparse
import json
import os
import sys

import chess
import torch

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from tinychess.batching import encode_batch
from tinychess.checkpoint import load_model_only


@torch.no_grad()
def trace_position(model, board, steps=8, adaptive=False, topk=5):
    pb = encode_batch([board], [board.peek() if board.move_stack else None])
    out = model(**pb.model_args(), steps=steps, adaptive=adaptive, trace=True)
    traj = out["trajectory"]
    coll = out.get("collect", {}) or {}
    moves = pb.real_moves[0]

    rec = {
        "fen": board.fen(),
        "side_to_move": "white" if board.turn else "black",
        "n_legal": int(pb.n_legal[0]),
        "params": model.param_report()["total"],
        "active_params": model.param_report()["active"],
        "adaptive": adaptive,
        "n_steps_used": float(out["n_steps"][0]),
        "value": float(out["value"][0, 0]),
        "steps": [],
    }

    prev = None
    for t, h in enumerate(traj):
        h0 = h[0]
        s = {"t": t,
             "h_norm": float(h0.norm()),
             "h_mean": float(h0.mean()),
             "h_std": float(h0.std()),
             "active_units": int((h0.abs() > 1e-3).float().mean(0).gt(0.5).sum()),
             }
        if prev is not None:
            d = h0 - prev
            s["delta_norm"] = float(d.norm())
            s["delta_rel"] = float(d.norm() / prev.norm().clamp_min(1e-6))
            s["most_changed_squares"] = [
                chess.square_name(i) for i in d.norm(dim=-1).topk(4).indices.tolist()]
        prev = h0
        i = t - 1
        if 0 <= i < len(coll.get("router", [])):
            w = coll["router"][i][0]
            s["route"] = {op: round(float(v), 3)
                          for op, v in zip(model.core.ops, w.tolist())}
        if 0 <= i < len(coll.get("halt", [])):
            s["p_halt"] = round(float(coll["halt"][i][0]), 4)
        if 0 <= i < len(coll.get("mem_attn", [])):
            ma = coll["mem_attn"][i][0]
            s["memory_attention"] = [round(float(v), 3) for v in ma.mean(0).tolist()]
        rec["steps"].append(s)

    if out.get("memory") is not None:
        m = out["memory"][0]
        rec["memory_final"] = {"n_slots": int(m.shape[0]),
                               "slot_norms": [round(float(v), 3) for v in m.norm(dim=-1)]}

    # per-refinement-step move ranking
    if "refine_logits" in out:
        rec["refinement"] = []
        for i, lg in enumerate(out["refine_logits"]):
            if lg is None:
                continue
            p = torch.softmax(lg[0], -1)
            k = min(topk, len(moves))
            top = p.topk(k)
            rec["refinement"].append({
                "z_index": i,
                "top": [{"move": moves[j].uci(), "p": round(float(v), 4)}
                        for v, j in zip(top.values, top.indices)]})

    p = torch.softmax(out["logits"][0], -1)
    k = min(topk, len(moves))
    top = p.topk(k)
    rec["final_policy"] = [{"move": moves[j].uci(), "p": round(float(v), 4)}
                           for v, j in zip(top.values, top.indices)]
    rec["chosen_move"] = moves[int(p.argmax())].uci()
    rec["policy_entropy"] = float(-(p * p.clamp_min(1e-12).log()).sum())
    return rec


def human(rec) -> str:
    L = []
    L.append("=" * 72)
    L.append(f"Position: {rec['fen']}")
    L.append(f"  side to move: {rec['side_to_move']}   legal moves: {rec['n_legal']}")
    L.append(f"  params: {rec['params']:,} (active {rec['active_params']:,})   "
             f"steps used: {rec['n_steps_used']:.2f}   value: {rec['value']:+.3f}")
    L.append("=" * 72)
    for s in rec["steps"]:
        L.append(f"Step {s['t']}")
        L.append(f"  |h|={s['h_norm']:.3f}  mean={s['h_mean']:+.4f}  std={s['h_std']:.4f}"
                 f"  active units={s['active_units']}")
        if "delta_norm" in s:
            L.append(f"  Δh={s['delta_norm']:.4f} (rel {s['delta_rel']:.3f})  "
                     f"changed: {', '.join(s['most_changed_squares'])}")
        if "route" in s:
            L.append("  route: " + "  ".join(f"{k}={v:.2f}" for k, v in s["route"].items()))
        if "p_halt" in s:
            L.append(f"  p_halt: {s['p_halt']:.3f}")
        if "memory_attention" in s:
            L.append(f"  memory attn: {s['memory_attention']}")
    if "memory_final" in rec:
        L.append(f"\nMemory ({rec['memory_final']['n_slots']} slots) norms: "
                 f"{rec['memory_final']['slot_norms']}")
    if "refinement" in rec:
        L.append("\nMove-latent refinement:")
        for r in rec["refinement"]:
            tops = "  ".join(f"{d['move']}:{d['p']:.3f}" for d in r["top"])
            L.append(f"  z{r['z_index']}: {tops}")
    L.append("\nFinal policy:")
    for d in rec["final_policy"]:
        L.append(f"  {d['move']}  p={d['p']:.4f}")
    L.append(f"\nFinal move: {rec['chosen_move']}   entropy={rec['policy_entropy']:.3f}")
    return "\n".join(L)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--checkpoint", required=True)
    ap.add_argument("--fen", default=chess.STARTING_FEN)
    ap.add_argument("--steps", type=int, default=8)
    ap.add_argument("--adaptive", action="store_true")
    ap.add_argument("--format", choices=["human", "json", "both"], default="human")
    ap.add_argument("--out", default=None)
    a = ap.parse_args()

    model, cfg = load_model_only(a.checkpoint)
    rec = trace_position(model, chess.Board(a.fen), a.steps, a.adaptive)
    if a.format in ("human", "both"):
        print(human(rec))
    if a.format in ("json", "both"):
        js = json.dumps(rec, indent=2)
        if a.out:
            open(a.out, "w").write(js)
            print(f"\n[saved json] {a.out}")
        elif a.format == "json":
            print(js)


if __name__ == "__main__":
    main()