"""Boot-prompt readout through the merged Huginn r=16 lens: per virtual layer,
transport h_v to the target and rank the prompt-absent answer token (Italy)."""

import torch
from jlens.lens import JacobianLens
from jlens.hooks import ActivationRecorder
from huginn_lens import HuginnLensModel, load_huginn

hf, tok = load_huginn()
hf = hf.cuda()
model = HuginnLensModel(hf, tok, num_steps=16)
lens = JacobianLens.load("jacobian_lens_huginn-0125_r16.pt")
print(model, flush=True)

prompt = "Fact: The currency used in the country shaped like a boot is"
ids = model.encode(prompt)
pos = ids.shape[1] - 1
italy = tok.encode(" Italy")[-1]
print("italy token:", italy, repr(tok.decode([italy])), flush=True)

layers = list(range(model.n_layers - 1))
with ActivationRecorder(model.layers, at=layers) as rec, torch.no_grad():
    out = model.forward(ids)

print(f"{'v':>3} {'kind':<10} {'top1':<16} {'italy_rank':>10}", flush=True)
for v in layers:
    h = rec.activations[v][0, pos].float().cuda()
    J = lens.jacobians[v].float().cuda()
    logits = model.unembed((h @ J.T).to(torch.bfloat16))
    rank = int((logits > logits[italy]).sum().item()) + 1
    top1 = tok.decode([logits.argmax().item()])
    if v < model.n_prelude:
        kind = "prelude"
    elif v >= model.n_prelude + model.num_steps * model.n_core:
        kind = "coda"
    else:
        kind = f"step{(v - model.n_prelude) // model.n_core}.L{(v - model.n_prelude) % model.n_core}"
    print(f"{v:>3} {kind:<10} {top1!r:<16} {rank:>10}", flush=True)
print("PROBE_DONE", flush=True)