"""Single-token-per-step latent CoT: ONE autoregressive soft token carries a whole K-cell state.

Layout: [X: prompt + <think>\\n][Z: T tokens, one per CA step][Y: </think> + query + boxed(+ans)].
Notation (see AGENTS.md): the autoregressively-fed soft tokens are z_t in R^{d_vocab}, t=1..T; the
residual stream is h; the ground-truth state is c in R^K. Here d_vocab = K*10 (a per-cell digit
distribution) and one z_t carries the ENTIRE K-cell row — so a reasoning step is a SINGLE token, not
K positions (contrast latent_threads/markov.py).

MARKOV mask: z_1 attends the prompt (the initial row c_0); z_t (t>1) attends ONLY z_{t-1}; the answer
attends ONLY z_T. Unique path prompt -> z_1 -> ... -> z_T -> answer, so every step is load-bearing.

Feedback: z_{t+1} = per-cell softmax(head(h at z_t)) in R^{K x 10}; embedded via a learned codebook
C in R^{K x 10 x d} (init = the model's digit embeddings + small per-cell noise -> readable, the
cells start near the plain digit embeddings and separate during training). The fed input embedding is
sum_{k,v} z[k,v] * C[k,v]. Teacher forcing substitutes the GT one-hot row (scheduled sampling).
"""
from __future__ import annotations

import torch
import torch.nn.functional as F

from abstract_cot.masking import PAD as ROLE_PAD, X as ROLE_X, Y as ROLE_Y, Z as ROLE_Z
from latent_threads import tasks as LT
from model_organisms.envs.base import initial_prefix_ids


def _ids(tok, s):
    return tok(s, add_special_tokens=False)["input_ids"]


def build_single_batch(tok, task, probs, im_end, device, with_answer=True):
    """[X][Z: m single-token steps][Y]. Returns ids/roles tensors + z_starts, label_starts, lengths."""
    m = task.m
    rows = []
    for p in probs:
        x = initial_prefix_ids(tok, task.prompt(p)) + _ids(tok, "<think>\n")
        y_open = _ids(tok, "\n</think>\n\n" + task.query(p) + "\\boxed{")
        y = y_open + (([LT.digit_ids(tok)[task.answer(p)]] + _ids(tok, "}") + [im_end]) if with_answer else [])
        ids = x + [tok.pad_token_id] * m + y                  # m latent tokens, one per step
        roles = [ROLE_X] * len(x) + [ROLE_Z] * m + [ROLE_Y] * len(y)
        rows.append((ids, roles, len(x), len(x) + m + len(y_open)))
    Lmax = max(len(r[0]) for r in rows)
    input_ids = torch.full((len(rows), Lmax), tok.pad_token_id, device=device)
    roles = torch.full((len(rows), Lmax), ROLE_PAD, device=device)
    z_starts, label_starts, lengths = [], [], []
    for j, (ids, rl, zs, ls) in enumerate(rows):
        input_ids[j, : len(ids)] = torch.tensor(ids, device=device)
        roles[j, : len(rl)] = torch.tensor(rl, device=device)
        z_starts.append(zs); label_starts.append(ls); lengths.append(len(ids))
    return input_ids, roles, z_starts, label_starts, lengths


def build_single_mask(roles, z_starts, m, dtype=torch.bfloat16, ablate_first=False):
    """Additive [B,1,L,L] single-token Markov mask: z_1->X, z_t->z_{t-1}, Y->z_m.
    ablate_first: ALSO blind z_1 to the prompt (the load-bearing control -> chance)."""
    B, L = roles.shape
    dev = roles.device
    neg = torch.finfo(dtype).min
    idx = torch.arange(L, device=dev)
    allowed = (idx[None, :] <= idx[:, None])[None].expand(B, L, L).clone()         # causal
    for b in range(B):
        zs = z_starts[b]; zend = zs + m
        zr = (roles[b] == ROLE_Z).nonzero(as_tuple=True)[0]
        xk = (roles[b] == ROLE_X).nonzero(as_tuple=True)[0]
        yi = (roles[b] == ROLE_Y).nonzero(as_tuple=True)[0]
        allowed[b][zr[:, None], xk[None, :]] = False                              # Z: forbid X & all Z
        allowed[b][zr[:, None], zs:zend] = False
        for t in range(m):
            pt = zs + t
            allowed[b][pt, pt] = True                                             # self
            if t == 0:
                if not ablate_first:
                    allowed[b][pt, xk] = True                                     # z_1 reads the prompt
            else:
                allowed[b][pt, zs + t - 1] = True                                 # z_t <- z_{t-1} only
        if len(yi):
            allowed[b][yi[:, None], xk[None, :]] = False                          # Y: forbid X & all Z but z_m
            allowed[b][yi[:, None], zs:zend - 1] = False
    pad = (roles == ROLE_PAD)
    allowed &= ~pad[:, None, :]
    add = torch.zeros((B, L, L), dtype=dtype, device=dev)
    add.masked_fill_(~allowed, neg)
    eye = torch.eye(L, dtype=torch.bool, device=dev)[None].expand(B, L, L)
    add = torch.where(pad[:, :, None] & eye, torch.zeros((), dtype=dtype, device=dev), add)
    add.masked_fill_((idx[None, :] > idx[:, None])[None].expand(B, L, L), neg)    # re-apply causality
    return add[:, None]


def single_forward(model, head, codebook, q_emb, tok, task, probs, im_end, device, with_answer=True,
                   tf_prob=0.0, gt_rng=None, ablate_first=False):
    """m in-graph forwards. z_1 input = learned q_emb; z_{t+1} = per-cell softmax(head(h at z_t)) @ C.
    Returns (answer_logits, meta, (aux_logits[B,m,K,10], aux_gt[B,m,K]))."""
    K, m = task.K, task.m
    input_ids, roles, z_starts, label_starts, lengths = build_single_batch(
        tok, task, probs, im_end, device, with_answer)
    B, Lmax = input_ids.shape
    attn = build_single_mask(roles, z_starts, m, ablate_first=ablate_first)
    pos = torch.arange(Lmax, device=device)[None].expand(B, Lmax)
    E = model.get_input_embeddings()(input_ids)
    bidx = torch.arange(B, device=device)
    zs = torch.tensor(z_starts, device=device)
    E = E.clone(); E[bidx, zs] = q_emb                                            # z_1 input = query vector
    gt_rows = torch.tensor([task.step_states(p) for p in probs], device=device)   # [B, m, K]
    tf_mask = None
    if tf_prob > 0 and gt_rng is not None:
        tf_mask = (torch.rand(B, generator=gt_rng, device="cpu") < tf_prob).to(device)
    aux_logits, aux_gt, out = [], [], None
    for t in range(m):
        out = model(inputs_embeds=E, attention_mask=attn, position_ids=pos, output_hidden_states=True)
        h = out.hidden_states[-1][bidx, zs + t]                                   # residual at z_t  [B, d]
        logits = head(h).view(B, K, 10)                                          # per-cell digit logits
        aux_logits.append(logits); aux_gt.append(gt_rows[:, t, :])
        if t == m - 1:
            break
        z = torch.softmax(logits.float(), dim=-1)                                # soft token z_t  [B,K,10]
        fed = torch.einsum("bkv,kvd->bd", z, codebook.float())
        if tf_mask is not None:
            gt_oh = F.one_hot(gt_rows[:, t, :], 10).float()                      # [B,K,10]
            fed_gt = torch.einsum("bkv,kvd->bd", gt_oh, codebook.float())
            fed = torch.where(tf_mask[:, None], fed_gt, fed)
        E = E.clone(); E[bidx, zs + t + 1] = fed.to(E.dtype)
    return out.logits, (z_starts, label_starts, lengths, input_ids), (torch.stack(aux_logits, 1), torch.stack(aux_gt, 1))


@torch.no_grad()
def single_readout_acc(model, head, codebook, q_emb, tok, task, probs, im_end, device,
                       tf_prob=0.0, gt_rng=None, ablate_first=False):
    """Free-running accuracy for the queried cell, read from the per-cell HEAD at the final token z_T.
    (The answer is decoded from z_T's residual by the same head that drives the recurrence -- a single
    readout path, no separate LM-head route to bypass.)"""
    model.eval()
    m = task.m
    correct, bs = 0, 16
    for i in range(0, len(probs), bs):
        batch = probs[i : i + bs]
        _, _, (aux_logits, _) = single_forward(
            model, head, codebook, q_emb, tok, task, batch, im_end, device, with_answer=False,
            tf_prob=tf_prob, gt_rng=gt_rng, ablate_first=ablate_first)
        final = aux_logits[:, m - 1]                                              # head at z_T: [B, K, 10]
        correct += sum(int(final[j, batch[j].q].argmax().item() == task.answer(batch[j])) for j in range(len(batch)))
    model.train()
    return correct / len(probs)


@torch.no_grad()
def single_state_acc(model, head, codebook, q_emb, tok, task, probs, im_end, device, tf_prob=0.0, gt_rng=None):
    """Per-cell state decodability: does head(h at z_t) recover the WHOLE row c_t? (mean over t, cells)"""
    model.eval()
    correct = tot = 0
    bs = 16
    for i in range(0, len(probs), bs):
        batch = probs[i : i + bs]
        _, _, (aux_logits, aux_gt) = single_forward(
            model, head, codebook, q_emb, tok, task, batch, im_end, device, with_answer=False,
            tf_prob=tf_prob, gt_rng=gt_rng)
        pred = aux_logits.argmax(-1)                                             # [B,m,K]
        correct += int((pred == aux_gt).sum()); tot += aux_gt.numel()
    model.train()
    return correct / tot