"""Reference inference for the block-distilled Qwen3-4B + decode-side per-step router.

Layout:  [system] [block_1 ...] [block_n] [query].  Each context block attends only to itself + the
system prefix (block-diagonal), so its KV is context-independent and REUSABLE across requests. Every
block is prefixed with 4 sink tokens and carries a `summ`-token SUMMARY (its last content tokens). The
query/decode attends to blocks; a light router selects which blocks stay ACTIVE.

Serving paths (all keep the shared block-diagonal prefill; differ only at decode):
  - dense       : decode attends to ALL blocks (quality reference).
  - reuse-shift : SELECT-ONCE. Router picks top-k active blocks at generation start, RoPE-shifts them +
                  query to COMPACT positions, decodes reusing that compact KV. == static KV pruning.
  - reroute     : PER-STEP (the contribution). ALL block KV stays resident; every decode step re-selects
                  the top-k active blocks from the current token's summary-attention, re-RoPE-shifts the
                  [system + active-full + all-summaries + query + prior-gen] to compact positions, and
                  decodes one token. Necessary when the RELEVANT block MOVES across the generated
                  sequence (long-form / multi-doc), where select-once/pruning structurally cannot follow.

RoPE key-shift is exact (relative RoPE: R(Δ)·R(p)k = R(p+Δ)k), so cached post-RoPE keys are retargeted
to compact positions WITHOUT recompute. This file is the SPEC for a custom serving kernel (e.g. vLLM):
block-diagonal context KV store + resident summaries + per-step top-k gather + RoPE compact-shift.

Usage:
    python infer.py --model <path> [--router router.pt --k 3] [--path dense|reuse-shift|reroute]
"""
from __future__ import annotations
import argparse, torch, torch.nn as nn
from transformers import AutoModelForCausalLM, AutoTokenizer, DynamicCache
from transformers.models.qwen3.modeling_qwen3 import rotate_half

NEG = -1e9; N_SINKS = 4; SINK_TEXT = "\n"
SYS = "You are a helpful assistant. Answer the question using only the documents. Be concise."
DOCS = [
    "The Zorban Reactor at Vale Station reached full output in the year 3471.",
    "Mount Kilimanjaro, in Tanzania, is the highest mountain in Africa at 5,895 metres.",
    "The Great Wall of China stretches thousands of kilometres across northern China.",
    "Photosynthesis converts carbon dioxide and water into glucose and oxygen.",
]
QUESTION = "In what year did the Zorban Reactor reach full output?"


def build_prompt(tok, docs, question):
    """Qwen3 CHAT layout (matches router training). ids, block id (-1 sys / -2 query / >=0 block), sink flag."""
    sink_ids = tok(SINK_TEXT, add_special_tokens=False)["input_ids"] * N_SINKS
    ids, blk, snk = [], [], []
    def add(text, b):
        for x in tok(text, add_special_tokens=False)["input_ids"]:
            ids.append(x); blk.append(b); snk.append(False)
    add(f"<|im_start|>system\n{SYS}<|im_end|>\n<|im_start|>user\n", -1)
    for bi, d in enumerate(docs):
        for s in sink_ids: ids.append(s); blk.append(bi); snk.append(True)
        add(f"[Document {bi+1}] {d}\n", bi)
    q0 = len(ids)
    add(f"{question}<|im_end|>\n<|im_start|>assistant\n<think>\n\n</think>\n\n", -2)
    return ids, blk, snk, q0


def block_summaries(blk, snk, nb, S):
    cols = {b: [] for b in range(nb)}
    for j, (b, s) in enumerate(zip(blk, snk)):
        if b >= 0 and not s: cols[b].append(j)
    return {b: cols[b][-S:] for b in range(nb)}


def prefill_mask(blk, snk, q0, nb, n, dev, dtype):
    """block-diagonal context + query-sees-all, causal. [1,1,n,n] additive, cast to model dtype (SDPA-safe)."""
    b = torch.tensor(blk + [-2] * (n - len(blk)), device=dev)
    causal = torch.tril(torch.ones(n, n, dtype=torch.bool, device=dev))
    bi, bj = b.view(n, 1), b.view(1, n)
    allowed = ((bj == -1) | (bi == bj) | (bi == -2)) & causal
    return torch.where(allowed, 0.0, NEG).view(1, 1, n, n).to(dtype)


def rope_shift_keys(keys, delta_pos, rotary, dev):
    cos, sin = rotary(keys, delta_pos[None].to(dev).float())
    return keys * cos.unsqueeze(1) + rotate_half(keys) * sin.unsqueeze(1)


def block_feat(att, smry, nb):
    """att [L,H,T] (decode token's attention). Feature = mean attention to each block's summary cols -> [nb,L*H]."""
    return torch.stack([att[:, :, smry[b]].mean(-1) for b in range(nb)], -1).permute(2, 0, 1).reshape(nb, -1)


@torch.no_grad()
def generate(model, tok, docs, question, router=None, k=3, max_new=32, summ=8, mode="dense"):
    """mode: dense | reuse-shift (select-once) | reroute (per-step). Router required for the sparse modes."""
    dev = model.device
    ids, blk, snk, q0 = build_prompt(tok, docs, question)
    n = len(ids); nb = max(blk) + 1
    smry = block_summaries(blk, snk, nb, summ); lmset = set(c for b in range(nb) for c in smry[b])
    rotary = model.model.rotary_emb
    out = model(input_ids=torch.tensor([ids], device=dev),
                attention_mask=prefill_mask(blk, snk, q0, nb, n, dev, model.dtype),
                use_cache=True, output_attentions=(mode != "dense"))
    pkv = out.past_key_values; nxt = int(out.logits[0, -1].argmax()); gen = [nxt]

    if mode == "dense":
        cur = n
        for _ in range(max_new - 1):
            if nxt == tok.eos_token_id: break
            cur += 1
            out = model(input_ids=torch.tensor([[nxt]], device=dev), attention_mask=torch.ones((1, cur), device=dev),
                        past_key_values=pkv, use_cache=True)
            nxt = int(out.logits[0, -1].argmax()); gen.append(nxt)
        return tok.decode(gen, skip_special_tokens=True), None

    # router pick at generation start (query-boundary attention to summaries)
    att0 = torch.stack(out.attentions, 0)[:, 0, :, -1, :].float()
    topk0 = set(router(block_feat(att0, smry, nb)).squeeze(-1).argsort(descending=True)[:k].tolist())

    if mode == "reuse-shift":                       # SELECT-ONCE: drop inactive KV, compact-shift, decode
        keep = [j for j, b in enumerate(blk) if b < 0 or b in topk0]
        old = torch.tensor(keep, device=dev); K = len(keep); delta = torch.arange(K, device=dev) - old
        for layer in pkv.layers:
            layer.keys = rope_shift_keys(layer.keys[:, :, old, :], delta, rotary, dev).contiguous()
            layer.values = layer.values[:, :, old, :].contiguous()
        cur = K
        for _ in range(max_new - 1):
            if nxt == tok.eos_token_id: break
            out = model(input_ids=torch.tensor([[nxt]], device=dev), position_ids=torch.tensor([[cur]], device=dev),
                        cache_position=torch.tensor([K + len(gen) - 1], device=dev), past_key_values=pkv, use_cache=True)
            nxt = int(out.logits[0, -1].argmax()); gen.append(nxt); cur += 1
        return tok.decode(gen, skip_special_tokens=True), sorted(topk0)

    # mode == "reroute": PER-STEP. keep ALL KV resident; re-select + re-compact every step.
    K0 = [l.keys for l in pkv.layers]; V0 = [l.values for l in pkv.layers]; L = len(K0)
    active = set(range(nb)); g = 0; Kg = [None] * L; Vg = [None] * L
    def keep_col(j, b): return (b < 0) or (b in active) or (j in lmset)
    for step in range(max_new - 1):
        kp = [j for j in range(n) if keep_col(j, blk[j])]
        pos_in = {j: i for i, j in enumerate(kp)}
        sidx = [[pos_in[c] for c in smry[b]] for b in range(nb)]
        Kt = len(kp) + g
        old = torch.tensor(kp + [n + i for i in range(g)], device=dev); delta = torch.arange(Kt, device=dev) - old
        cache = DynamicCache()
        for i in range(L):
            kk = K0[i][:, :, kp, :]; vv = V0[i][:, :, kp, :]
            if g: kk = torch.cat([kk, Kg[i]], 2); vv = torch.cat([vv, Vg[i]], 2)
            cache.update(rope_shift_keys(kk, delta, rotary, dev), vv, i)
        out = model(input_ids=torch.tensor([[nxt]], device=dev), position_ids=torch.tensor([[Kt]], device=dev),
                    cache_position=torch.tensor([Kt], device=dev), past_key_values=cache, use_cache=True,
                    output_attentions=True)
        nxt2 = int(out.logits[0, -1].argmax())
        att = torch.stack(out.attentions, 0)[:, 0, :, 0, :].float()
        active = set(router(block_feat(att, {b: sidx[b] for b in range(nb)}, nb)).squeeze(-1).argsort(descending=True)[:k].tolist())
        dstore = torch.tensor([(n + g) - Kt], device=dev)
        for i in range(L):
            Kg[i] = rope_shift_keys(cache.layers[i].keys[:, :, -1:, :], dstore, rotary, dev) if Kg[i] is None \
                    else torch.cat([Kg[i], rope_shift_keys(cache.layers[i].keys[:, :, -1:, :], dstore, rotary, dev)], 2)
            Vg[i] = cache.layers[i].values[:, :, -1:, :] if Vg[i] is None \
                    else torch.cat([Vg[i], cache.layers[i].values[:, :, -1:, :]], 2)
        g += 1; gen.append(nxt2); nxt = nxt2
        if nxt == tok.eos_token_id: break
    return tok.decode(gen, skip_special_tokens=True), None


def load_router(path, dev):
    ck = torch.load(path, map_location=dev)
    LH = ck["in_dim"]; arch = ck.get("arch", "linear"); summ = ck.get("summary_tokens", 8)
    if arch == "mlp":
        r = nn.Sequential(nn.Linear(LH, 256), nn.GELU(), nn.Dropout(0.0), nn.Linear(256, 1))
    else:
        r = nn.Linear(LH, 1)
    r.load_state_dict(ck["state_dict"]); return r.to(dev).eval(), summ


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--model", required=True)
    ap.add_argument("--router", default=None); ap.add_argument("--k", type=int, default=3)
    ap.add_argument("--max-new", type=int, default=32)
    ap.add_argument("--path", default="reroute", choices=["dense", "reuse-shift", "reroute"])
    args = ap.parse_args()
    tok = AutoTokenizer.from_pretrained(args.model)
    model = AutoModelForCausalLM.from_pretrained(args.model, dtype=torch.bfloat16,
              attn_implementation="eager", device_map="cuda").eval()
    router = summ = None
    if args.router: router, summ = load_router(args.router, "cuda")
    ids, blk, _, q0 = build_prompt(tok, DOCS, QUESTION)
    print(f"blocks={max(blk)+1}  context_tokens={q0}  path={args.path}  question={QUESTION!r}")
    dense, _ = generate(model, tok, DOCS, QUESTION, max_new=args.max_new, mode="dense")
    print("dense :", repr(dense.strip()))
    if router is not None and args.path != "dense":
        s, act = generate(model, tok, DOCS, QUESTION, router, args.k, args.max_new, summ, args.path)
        print(f"{args.path} (top-{args.k}, active={act}) :", repr(s.strip()))


if __name__ == "__main__":
    main()