"""Deliberation-trace stage 1: per-recurrence J-lens readout over a riddle set
with known plausible-but-wrong candidates (full 16-lens tiling).

Per riddle, one recorded forward. At every virtual layer: rank of the
expected answer and each trap candidate, projection of the state onto each
token's lens direction (plus fixed control tokens), and the top-8 decoded
tokens. Plus a 12-token greedy continuation for the model's actual answer.
"""

import json

import torch

import ws_common as W

torch.set_grad_enabled(False)

ctx = W.load_ctx("huginn", full_depth=True, huginn_full_tiling=True)
V = 67  # virtual layers with a covering lens (0..66)

RIDDLES = [
    ("boot_currency", "Fact: The currency used in the country shaped like a boot is the",
     "euro", ["lira", "Italy", "Italian"]),
    ("eiffel_pre2002", "The currency used in the country of the Eiffel Tower before 2002 was the",
     "franc", ["euro", "France"]),
    ("sushi_currency", "The currency used in the country famous for sushi is the",
     "yen", ["yuan", "won"]),
    ("kangaroo_capital", "The capital of the country where kangaroos live is",
     "Canberra", ["Sydney", "Melbourne"]),
    ("maple_capital", "The capital of the country famous for maple syrup is",
     "Ottawa", ["Toronto", "Montreal"]),
    ("liberty_capital", "The capital of the country home to the Statue of Liberty is",
     "Washington", ["New", "York", "Paris"]),
    ("largest_planet", "The largest planet in the solar system is",
     "Jupiter", ["Saturn", "Neptune"]),
    ("closest_planet", "The planet closest to the sun is",
     "Mercury", ["Venus", "Mars"]),
    ("fastest_animal", "The fastest land animal is the",
     "cheetah", ["lion", "leopard"]),
    ("jungle_king", "The animal known as the king of the jungle is the",
     "lion", ["tiger", "elephant"]),
    ("tallest_mountain", "The tallest mountain in the world is Mount",
     "Everest", ["Fuji", "Kilimanjaro"]),
    ("relativity", "The scientist who developed the theory of relativity was",
     "Einstein", ["Newton", "Darwin"]),
]

CONTROLS = ["table", "music", "seven", "window", "doctor", "river",
            "yellow", "engine"]


def greedy(prompt, n=12):
    ids = W.encode(ctx, prompt)
    outs = []
    for _ in range(n):
        logits, _ = W.forward(ctx, ids, layers=[])
        t = int(logits[-1].argmax())
        outs.append(t)
        ids = torch.cat([ids, torch.tensor([[t]], device=ids.device)], dim=1)
    return ctx.tok.decode(outs)


out = {"riddles": []}
ctl_toks = [W.first_token(ctx, w) for w in CONTROLS]
for name, prompt, answer, traps in RIDDLES:
    toks = {answer: W.first_token(ctx, answer)}
    for t in traps:
        toks[t] = W.first_token(ctx, t)
    logits, acts = W.forward(ctx, W.encode(ctx, prompt), layers=list(range(V)))
    rec = {"name": name, "prompt": prompt, "answer": answer, "traps": traps,
           "argmax": ctx.tok.decode([int(logits[-1].argmax())]),
           "gen": greedy(prompt),
           "ranks": {w: [] for w in toks}, "proj": {w: [] for w in toks},
           "ctl_proj": [], "top8": []}
    for v in range(V):
        h = acts[v][-1]
        lg = ctx.model.unembed(h @ ctx.Jt[v]).float()
        for w, t in toks.items():
            rec["ranks"][w].append(int((lg > lg[t]).sum().item()) + 1)
            d = W.direction(ctx, v, t)
            rec["proj"][w].append(float(h @ d))
        rec["ctl_proj"].append([float(h @ W.direction(ctx, v, t))
                                for t in ctl_toks])
        rec["top8"].append([ctx.tok.decode([t]) for t in
                            lg.topk(8).indices.tolist()])
    out["riddles"].append(rec)
    print(f"done {name}: argmax={rec['argmax']!r} gen={rec['gen'][:40]!r}",
          flush=True)

json.dump(out, open("delib_trace.json", "w"))
print("wrote delib_trace.json", flush=True)
print("DONE_DELIB", flush=True)