"""Inference-only runtime for the trace transformer (tiers 3-6).

Ships in the submission. Contains ONLY: token constants, base-256 limb
conversion (base conversion is explicitly allowed representation work),
prompt construction, and a batched KV-cache greedy decode loop that collects
the model-emitted answer digits after the ANS marker.

No arithmetic on input values happens here: Python only selects and collects
model-emitted token ids. A fixed, value-independent selection (slice after
ANS) maps the model's emission to the returned digit list; garbage emissions
produce garbage (or sentinel) answers.
"""

from __future__ import annotations

import torch

# Vocabulary (must match training; see model_trace.py)
PAD, SEP, EQ, ANS, EOS = 256, 257, 258, 259, 260

# tier -> (min p bits, max p bits, limbs per field)
TRACE_TIERS = {
    3: (9, 16, 2),
    4: (17, 32, 4),
    5: (33, 64, 8),
    6: (65, 128, 16),
}


def tier_of_bits(bits: int) -> int | None:
    for t, (lo, hi, _nl) in TRACE_TIERS.items():
        if lo <= bits <= hi:
            return t
    return None


def expected_trace_tokens(n_limbs: int) -> int:
    # v3 grammar: 12L^2 + 38L + 23 (see dev/trace/grammar.py).
    # Pure length estimate as a function of the tier's limb count — used
    # only to size the decode loop / time budget. No input-value arithmetic.
    return 12 * n_limbs * n_limbs + 38 * n_limbs + 23


def int_to_limbs(x: int, n_limbs: int) -> list[int]:
    """Base-256 digits, MSB-first, zero-padded to n_limbs."""
    out = [0] * n_limbs
    i = n_limbs - 1
    while x > 0 and i >= 0:
        out[i] = x & 0xFF
        x >>= 8
        i -= 1
    return out


def prompt_tokens(a_red: int, b_red: int, p: int, n_limbs: int) -> list[int]:
    return (
        int_to_limbs(p, n_limbs) + [SEP]
        + int_to_limbs(a_red, n_limbs) + [SEP]
        + int_to_limbs(b_red, n_limbs) + [EQ]
    )


@torch.no_grad()
def decode_batch(model, problems, device, sentinel, max_new_factor: float = 1.3):
    """Batched greedy decode with KV cache.

    problems: list of (a_red, b_red, p, n_limbs) — all with the SAME n_limbs.
    Returns list of digit lists (n_limbs ints in [0,255], MSB-first) or
    `sentinel` (copied) where decoding failed to produce a full answer.
    """
    if not problems:
        return []
    model.eval()
    model.clear_cache()
    B = len(problems)
    n_limbs = problems[0][3]
    max_new = int(expected_trace_tokens(n_limbs) * max_new_factor) + 8

    prompts = [prompt_tokens(a, b, p, nl) for (a, b, p, nl) in problems]
    # Fixed limb widths -> all prompts in a group have identical length.
    ctx = torch.tensor(prompts, dtype=torch.long, device=device)  # [B, T_p]

    logits = model(ctx, use_cache=True)            # primes the KV cache
    next_toks = logits[:, -1, :].argmax(-1)        # [B]

    ans_seen = [False] * B
    ans_digits: list[list[int]] = [[] for _ in range(B)]
    done = [False] * B

    for _ in range(max_new):
        all_done = True
        toks = next_toks.tolist()
        for i in range(B):
            if done[i]:
                continue
            tok = toks[i]
            if tok == EOS:
                done[i] = True
            elif tok == ANS:
                ans_seen[i] = True
            elif ans_seen[i] and tok < 256:
                ans_digits[i].append(tok)
                if len(ans_digits[i]) == n_limbs:
                    done[i] = True
            if not done[i]:
                all_done = False
        if all_done:
            break
        inp = next_toks.clone()
        step_logits = model(inp.unsqueeze(1), use_cache=True)  # [B, 1, V]
        next_toks = step_logits[:, -1, :].argmax(-1)

    model.clear_cache()
    return [
        d if len(d) == n_limbs else list(sentinel)
        for d in ans_digits
    ]