"""Experiment 1: J-lens readouts.

(a) Two-hop reasoning: unspoken intermediate 'spider' visible mid-network.
(b) Multilingual: English intermediates during a Chinese task.
(c) Directed modulation: 'hold citrus in mind' while copying unrelated text.
(d) Layer profile: where in depth the lens carries abstract content.
"""

import os, sys
from pathlib import Path

import torch

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
from jlens.core import (JLens, chat_ids, collect_residuals,
                        generate_with_residuals, load_model)

RES = Path(os.environ.get("JLENS_RESULTS", "results"))


def show_table(jl, tok, ids, layers, positions, topk=6, header=""):
    """Print top J-lens tokens at (layer, position) grid."""
    hs = collect_residuals(jl.model, ids)
    toks = [tok.decode([t]) for t in ids[0].tolist()]
    print(f"\n=== {header} ===")
    print("positions:", {p: repr(toks[p]) for p in positions})
    for l in layers:
        row = []
        for p in positions:
            idx, prob = jl.read(hs[l, p], l, topk=topk)
            row.append(" ".join(
                (tok.decode([i]).strip() or "·") for i in idx.tolist()[:topk]))
        print(f"L{l:>2} | " + " || ".join(row))


def concept_heatmap(jl, tok, ids, words, tag, hs=None):
    """P(concept tokens) by (layer, position); save tensor + print peak."""
    if hs is None:
        hs = collect_residuals(jl.model, ids)
    tids = []
    for w in words:
        for v in (w, " " + w, w.capitalize(), " " + w.capitalize()):
            e = tok.encode(v, add_special_tokens=False)
            if len(e) == 1:
                tids.append(e[0])
    tids = sorted(set(tids))
    L, T, _ = hs.shape
    out = torch.zeros(L, T)
    for l in range(L):
        logits = jl._readout(hs[l].float() @ jl.Jbar[l].T)
        probs = torch.softmax(logits.float(), dim=-1)
        out[l] = probs[:, tids].sum(-1).cpu()
    peak = out.max().item()
    lmax, pmax = divmod(out.argmax().item(), T)
    toks = [tok.decode([t]) for t in ids[0].tolist()] if ids is not None else None
    print(f"[{tag}] words={words} peak P={peak:.3f} at layer {lmax}, "
          f"pos {pmax}" + (f" ({toks[pmax]!r})" if toks and pmax < len(toks) else ""))
    RES.mkdir(exist_ok=True)
    torch.save({"map": out, "words": words,
                "tokens": [tok.decode([t]) for t in ids[0].tolist()] if ids is not None else None},
               RES / f"heat_{tag}.pt")
    return out


def main():
    model, tok = load_model(dtype=torch.bfloat16)
    ck = torch.load(sys.argv[1] if len(sys.argv) > 1 else "results/jbar.pt",
                    map_location="cuda")
    jl = JLens(model, tok, ck["Jbar"].cuda())
    print(f"Jbar from {ck.get('n_prompts')} prompts")
    L = len(model.model.language_model.layers)
    mid = range(max(2, L // 6), L - 4, max(2, L // 16))

    # (a) two-hop: spider
    ids = chat_ids(tok, "The animal that spins webs has how many legs? "
                        "Answer with just a number.")
    print("answer:", jl.generate(ids, 5))
    show_table(jl, tok, ids, mid, list(range(ids.shape[1] - 10, ids.shape[1])),
               header="two-hop spider (last 10 positions)")
    concept_heatmap(jl, tok, ids, ["spider", "spiders"], "spider")
    concept_heatmap(jl, tok, ids, ["eight", "8"], "eight")

    # (b) multilingual: Chinese antonym of small -> big
    ids = chat_ids(tok, "小的反义词是什么？只用一个字回答。")
    print("\nanswer:", jl.generate(ids, 5))
    show_table(jl, tok, ids, mid, list(range(ids.shape[1] - 8, ids.shape[1])),
               header="Chinese antonym of 小")
    concept_heatmap(jl, tok, ids, ["big", "large", "bigger"], "big_en")

    # (c) directed modulation: think of citrus while copying text
    copy_text = "The committee will meet on Thursday to review the budget."
    for cond, instr in [
        ("citrus", "While you copy the text, silently think about citrus "
                   "fruits the entire time. Copy this text exactly, output "
                   f"nothing else: \"{copy_text}\""),
        ("control", f"Copy this text exactly, output nothing else: \"{copy_text}\""),
    ]:
        ids = chat_ids(tok, instr)
        text, all_ids, hs = generate_with_residuals(model, tok, ids, 30)
        print(f"\n[{cond}] output: {text!r}")
        # only look at generated positions
        gen0 = ids.shape[1]
        heat = concept_heatmap(jl, tok, all_ids.unsqueeze(0),
                               ["citrus", "lemon", "orange", "lime"],
                               f"citrus_{cond}", hs=hs)
        print(f"  citrus P over generated positions: mean "
              f"{heat[:, gen0:].mean():.4f} max {heat[:, gen0:].max():.4f}")

    # (d) layer profile: lens entropy + top-token agreement with input/output
    ids = chat_ids(tok, "Write one sentence about the ocean.")
    text, all_ids, hs = generate_with_residuals(model, tok, ids, 20)
    L, T, _ = hs.shape
    ent, agree_in, agree_next = [], [], []
    for l in range(L):
        logits = jl._readout(hs[l].float() @ jl.Jbar[l].T)
        probs = torch.softmax(logits.float(), dim=-1)
        e = -(probs * probs.clamp_min(1e-12).log()).sum(-1).mean()
        top = probs.argmax(-1)
        ent.append(e.item())
        agree_in.append((top[:-1] == all_ids[:len(top) - 1].cuda()).float().mean().item())
        agree_next.append((top[:-1] == all_ids[1:len(top)].cuda()).float().mean().item())
    print("\nlayer | lens entropy | top==current tok | top==next tok")
    for l in range(0, L, 2):
        print(f"L{l:>2} | {ent[l]:8.2f} | {agree_in[l]:.2f} | {agree_next[l]:.2f}")
    torch.save({"entropy": ent, "agree_in": agree_in, "agree_next": agree_next},
               RES / "layer_profile.pt")


if __name__ == "__main__":
    main()