"""Companion to probe_discount.py: lens table over the PROMPT positions
after k=2 prefill settling (arm A adapter) — what the workspace holds
about the question before any pause/answer compute begins."""
import os
import sys
from pathlib import Path

import torch

from carry_common import _pos_ids, prompt_prefill
from loop_common import BandLooper, MergeAdapter, chat_prompt, DIRECT_SUFFIX

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
from jlens.core import JLens, ResidualCapture, load_model  # noqa: E402

OUT = Path(os.environ.get("LOOP_OUT",
    Path(__file__).resolve().parent.parent / "results-loop"))
Q = "A shirt costs $80 and is on sale for 15% off. How much does it cost?"
K = 2
TAG = "discount"
if "--idx" in sys.argv:
    import json
    _i = int(sys.argv[sys.argv.index("--idx") + 1])
    _star = {it["idx"]: it for it in json.load(open(OUT / "star_data.json"))
             if it["split"] == "test"}
    Q = _star[_i]["question"]
    TAG = f"gsm{_i}"

model, tok = load_model(dtype=torch.bfloat16)
looper = BandLooper(model)
band = list(range(looper.l0, looper.l1 + 1))
adapter = MergeAdapter(d=model.config.get_text_config().hidden_size).cuda()
adapter.load_state_dict(torch.load(OUT / "adapter_carrycot_e400.pt",
                                   map_location="cuda"))
jbar = torch.load(Path(__file__).resolve().parent.parent / "results/jbar.pt",
                  map_location="cuda")["Jbar"]
jl = JLens(model, tok, jbar)

prompt = chat_prompt(tok, Q, DIRECT_SUFFIX)
ids = tok(prompt, add_special_tokens=False,
          return_tensors="pt")["input_ids"].cuda()
mask = torch.ones_like(ids)

with torch.no_grad():
    calls, _ = looper.capture(ids, mask, logits_to_keep=1,
                              position_ids=_pos_ids(mask))
    e = looper._hin[looper.l0]
    pmask = torch.ones_like(ids, dtype=torch.bool)
    S, X = prompt_prefill(looper, adapter, e, calls, pmask, K)
    with ResidualCapture(model, layers=band) as rc:
        looper.band(X, calls)
    acts = {l: rc.acts[l][0] for l in band}
    # k=0 reference: plain forward, no adapter anywhere
    with ResidualCapture(model, layers=band) as rc0:
        model(input_ids=ids, attention_mask=mask, use_cache=False)
    acts0 = {l: rc0.acts[l][0] for l in band}

table = {}
for l in band:
    idx, p = jl.read(acts[l], l, topk=5)
    table[l] = (idx.cpu(), p.cpu())
idx0, p0 = jl.read(acts0[30], 30, topk=5)
toks = [repr(tok.decode([t])) for t in ids[0].tolist()]
torch.save({"question": Q, "k": K, "poslab": toks, "band": band,
            "table": table, "base30": (idx0.cpu(), p0.cpu())},
           OUT / f"probe_{TAG}_prompt.pt")
show = [14, 18, 22, 26, 30]
print(f"=== prompt board after k={K} settle (arm A)  ||  k=0 base ===")
for i, lab in enumerate(toks):
    row, row0 = [], []
    for l in show:
        idx, _ = jl.read(acts[l][i][None], l, topk=3)
        row.append("/".join(repr(tok.decode([t]))[1:-1]
                            for t in idx[0].tolist()))
    idx0, _ = jl.read(acts0[30][i][None], 30, topk=3)
    base30 = "/".join(repr(tok.decode([t]))[1:-1] for t in idx0[0].tolist())
    print(f"{i:3d} {lab:14s} " + "  ".join(
        f"L{l}:{r}" for l, r in zip(show, row)) + f"   || base-L30:{base30}",
        flush=True)