"""J-lens sharpening battery on MBPP prompts: does the trained loop
concentrate the workspace readout while reading the problem?

For N test prompts, at each k: J-lens distribution at L30, last prompt
position -> entropy + top-1 probability. Trained vs untrained adapter.
"""

import json
import sys
from pathlib import Path

import torch

from loop_common import BandLooper, MergeAdapter
from prep_mbpp import DIRECT_SUFFIX, mbpp_prompt

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

OUT = Path(__file__).resolve().parent.parent / "results-loop"
ROOT = OUT.parent
N_ITEMS = 20
KS = (0, 1, 2, 4)


@torch.no_grad()
def battery(looper, adapter, lens, tok, items):
    rows = []
    for it in items:
        ids = tok(mbpp_prompt(tok, it, DIRECT_SUFFIX), return_tensors="pt",
                  add_special_tokens=False)["input_ids"].cuda()
        calls, _ = looper.capture(ids, logits_to_keep=1)
        e = looper._hin[looper.l0]
        s = looper.band(e, calls)
        state = {0: s}
        for k in range(1, max(KS) + 1):
            s = looper.band(adapter(e, s), calls)
            state[k] = s
        for k in KS:
            probs = torch.softmax(lens._readout(
                state[k][0, -1].float() @ lens.Jbar[looper.l1].T).float(), -1)
            ent = -(probs * (probs + 1e-12).log()).sum().item()
            rows.append({"task_id": it["task_id"], "k": k, "entropy": ent,
                         "top1": probs.max().item()})
    return rows


def main():
    model, tok = load_model(dtype=torch.bfloat16)
    looper = BandLooper(model)
    jbar = torch.load(ROOT / "results" / "jbar.pt", map_location="cuda")
    lens = JLens(model, tok, jbar["Jbar"].float())
    items = [it for it in json.load(open(OUT / "mbpp_data.json"))
             if it["split"] == "test"][:N_ITEMS]

    res = {}
    for tag, path in (("untrained", None),
                      ("trained", OUT / "adapter_code_e399.pt")):
        adapter = MergeAdapter().cuda()
        if path:
            adapter.load_state_dict(torch.load(path, map_location="cuda"))
        rows = battery(looper, adapter, lens, tok, items)
        res[tag] = rows
        for k in KS:
            sel = [r for r in rows if r["k"] == k]
            ent = sum(r["entropy"] for r in sel) / len(sel)
            top = sum(r["top1"] for r in sel) / len(sel)
            print(f"[{tag}] k={k}: mean_entropy={ent:.3f}  mean_top1={top:.3f}",
                  flush=True)
    json.dump(res, open(OUT / "lens_battery.json", "w"), indent=1)
    print("wrote", OUT / "lens_battery.json")


if __name__ == "__main__":
    main()