"""Design C machinery: k-loop the prompt, then per-token carry generation.

Structure per sequence:  [prompt] [p pause tokens] [answer]
  1. prefill: prompt positions looped k times through the merge (settled plan)
  2. carry scan: each pause/answer position's band input is
         x_t = merge(e_t, s_{t-1})
     where s_{t-1} is the *previous position's* final band output — the
     latent thought chain runs through the pause positions' workspace states
     (token content of pauses is inert: <unused0>).
  3. suffix -> logits; CE on answer tokens only.

The scan is sequential per position (RNN-style through the band); each step
re-runs the band on the full sequence — causality keeps settled positions
stable, provisional later positions are recomputed next step anyway.
"""

import torch
from torch.utils.checkpoint import checkpoint

PAUSE_ID = 6  # <unused0>


def _pos_ids(mask):
    return (mask.cumsum(-1) - 1).clamp(min=0)


def prompt_prefill(looper, adapter, e, calls, prompt_mask, k):
    """k merge->band loops over prompt positions. Returns (S, X)."""
    with torch.no_grad():
        s = looper.band(e, calls)
    x = e
    for _ in range(k):
        x = torch.where(prompt_mask[..., None], adapter(e, s), e)
        s = looper.band(x, calls)
    return s, x


def carry_steps(looper, adapter, e, calls, S, X, step_updates,
                use_checkpoint=False, iter_states=None):
    """Sequential scan. step_updates: list of (row_idx, pos, [inplace])
    index tensors. inplace=True re-iterates the SAME position, seeding
    from its own previous band output (internal band looping) instead of
    its left neighbor. iter_states: optional list — every inplace
    update's fresh band output at pos is appended (for lens losses)."""
    for upd in step_updates:
        rows, pos = upd[0], upd[1]
        inplace = len(upd) > 2 and upd[2]
        if rows.numel() == 0:
            continue
        seed = S[rows, pos] if inplace else S[rows, pos - 1]
        x_new = adapter(e[rows, pos], seed)
        X = X.clone()
        X[rows, pos] = x_new.to(X.dtype)
        S = (checkpoint(lambda X_: looper.band(X_, calls), X,
                        use_reentrant=False) if use_checkpoint
             else looper.band(X, calls))
        if inplace and iter_states is not None:
            iter_states.append(S[rows, pos])
    return S, X


def build_step_updates(prompt_lens, total_lens, device):
    """For right-padded batches: step j updates row b at prompt_lens[b]+j."""
    max_span = int((total_lens - prompt_lens).max())
    updates = []
    for j in range(max_span):
        pos = prompt_lens + j
        sel = pos < total_lens
        rows = torch.nonzero(sel, as_tuple=False).squeeze(-1)
        updates.append((rows.to(device), pos[sel].to(device)))
    return updates


def splice_inner_iters(updates, inner_iters, inner_at, prompt_lens, dev, B):
    """Insert inner_iters in-place band iterations at the (batch-uniform
    offset) anchor position, right after the scan first settles it."""
    j_anchor = int((inner_at - prompt_lens).max())
    rows = torch.arange(B, device=dev)
    inner = [(rows, inner_at.to(dev), True)] * inner_iters
    if j_anchor < 0:
        # anchor is the last PROMPT position (no pauses at all):
        # iterate there before the scan enters the visible tokens
        return inner + updates
    out = []
    for j, u in enumerate(updates):
        out.append(u)
        if j == j_anchor:
            out += inner
    return out


def sym_iterate(looper, adapter, proj, e, calls, S, X, rows, anchor, m,
                lens_fn, embed_w, sym_tf=None, start_id=None, topk=32,
                use_checkpoint=False, iter_states=None):
    """Item 32: discrete latent chain at the anchor. Each tick reads the
    previous anchor state through the lens, snaps it to a token
    (straight-through over top-k) or takes the teacher token (sym_tf:
    (B, m) ids, teacher forcing), and feeds that token's embedding back
    through a zero-init projector ALONGSIDE the analog carry:
        x_i = merge(e, s_{i-1}) + proj(E(sym))
    Tick 0 uses start_id (a newline: 'a step begins')."""
    for i in range(m):
        s_prev = S[rows, anchor]
        if sym_tf is not None:
            symb = embed_w[sym_tf[:, i]]
        elif i == 0:
            symb = embed_w[torch.full((rows.shape[0],), start_id,
                                      device=e.device)]
        else:
            logits = lens_fn(s_prev).float()
            p, idx = torch.softmax(logits, -1).topk(topk, dim=-1)
            p = p / p.sum(-1, keepdim=True)
            soft = (p.unsqueeze(-1) * embed_w[idx].float()).sum(-2)
            hard = embed_w[idx[:, 0]].float()
            symb = hard + soft - soft.detach()
        x_new = (adapter(e[rows, anchor], s_prev).float()
                 + proj(symb.float()))
        X = X.clone()
        X[rows, anchor] = x_new.to(X.dtype)
        S = (checkpoint(lambda X_: looper.band(X_, calls), X,
                        use_reentrant=False) if use_checkpoint
             else looper.band(X, calls))
        if iter_states is not None:
            iter_states.append(S[rows, anchor])
    return S, X


def carry_logits(looper, adapter, input_ids, attention_mask, prompt_lens,
                 k, use_checkpoint=False, feedforward=False,
                 return_states=False, inner_iters=0, inner_at=None,
                 iter_states=None):
    """Teacher-forced design-C forward (right-padded batch).

    feedforward=True: pause-token control — same positions get the adapter as
    x_t = merge(e_t, e_t) in ONE band pass; no state propagates between
    positions. Isolates 'pause compute + weights' from the recurrence."""
    dev = input_ids.device
    calls, _ = looper.capture(input_ids, attention_mask, logits_to_keep=1)
    e = looper._hin[looper.l0].detach()
    ar = torch.arange(input_ids.shape[1], device=dev)
    prompt_mask = ar[None, :] < prompt_lens[:, None].to(dev)
    if feedforward:
        touched = attention_mask.bool()
        x = torch.where(touched[..., None], adapter(e, e), e)
        S = (checkpoint(lambda x_: looper.band(x_, calls), x,
                        use_reentrant=False) if use_checkpoint
             else looper.band(x, calls))
        out = looper.suffix_logits(S, calls)
        return (out, S) if return_states else out
    S, X = prompt_prefill(looper, adapter, e, calls, prompt_mask, k)
    total_lens = attention_mask.sum(-1)
    updates = build_step_updates(prompt_lens.to(dev), total_lens.to(dev), dev)
    if inner_iters and inner_at is not None:
        updates = splice_inner_iters(updates, inner_iters, inner_at,
                                     prompt_lens.to(dev), dev,
                                     input_ids.shape[0])
    S, X = carry_steps(looper, adapter, e, calls, S, X, updates,
                       use_checkpoint=use_checkpoint,
                       iter_states=iter_states)
    out = looper.suffix_logits(S, calls)
    return (out, S) if return_states else out


@torch.no_grad()
def generate_carry_c(looper, adapter, tok, input_ids, attention_mask,
                     k, p, max_new_tokens=10, feedforward=False,
                     inner_iters=0, kvmem=None, symchain=None):
    """Greedy design-C generation (left-padded batch, uniform positions).

    Appends p pause tokens, prefill-loops the prompt, carries through the
    pauses, then generates with per-token carry."""
    B = input_ids.shape[0]
    dev = input_ids.device
    pauses = torch.full((B, p), PAUSE_ID, dtype=torch.long, device=dev)
    ids = torch.cat([input_ids, pauses], 1)
    mask = torch.cat([attention_mask,
                      torch.ones_like(pauses)], 1)
    n_prompt = input_ids.shape[1]

    eos = {tok.eos_token_id, tok.convert_tokens_to_ids("<end_of_turn>")}
    done = torch.zeros(B, dtype=torch.bool, device=dev)
    X_store = None  # merged band inputs for settled positions
    for step in range(max_new_tokens + 1):  # step 0 = prefill + pause scan
        calls, _ = looper.capture(ids, mask, logits_to_keep=1,
                                  position_ids=_pos_ids(mask))
        e = looper._hin[looper.l0]
        if feedforward:
            x = torch.where(mask.bool()[..., None], adapter(e, e), e)
            S = looper.band(x, calls)
            X = x
        elif X_store is None:
            ar = torch.arange(ids.shape[1], device=dev)
            prompt_mask = (ar[None, :] < n_prompt) & mask.bool()
            S, X = prompt_prefill(looper, adapter, e, calls, prompt_mask, k)
            updates = [(torch.arange(B, device=dev),
                        torch.full((B,), n_prompt + j, device=dev,
                                   dtype=torch.long)) for j in range(p)]
            if symchain is not None and inner_iters:
                anchor_sc = torch.full((B,), n_prompt + p - 1, device=dev,
                                       dtype=torch.long)
                S, X = sym_iterate(
                    looper, adapter, symchain["proj"], e, calls, S, X,
                    torch.arange(B, device=dev), anchor_sc, inner_iters,
                    symchain["lens_fn"], symchain["embed_w"],
                    start_id=symchain["start_id"])
                updates = updates  # pauses (if any) already handled above
            elif inner_iters:
                anchor = torch.full((B,), n_prompt + p - 1, device=dev,
                                    dtype=torch.long)
                inner = [(torch.arange(B, device=dev), anchor, True)
                         for _ in range(inner_iters)]
                # p=0: iterate at the last prompt position, before any
                # visible token — no pause tokens involved
                updates = (inner + updates) if p == 0 else (updates + inner)
            itst = [] if (kvmem is not None and inner_iters) else None
            S, X = carry_steps(looper, adapter, e, calls, S, X, updates,
                               iter_states=itst)
            if kvmem is not None and itst:
                from kv_memory import arm_memory
                arm_memory(kvmem(torch.stack(itst, 1)))
        else:
            X = torch.cat([X_store, e[:, X_store.shape[1]:]], 1)
            rows = torch.arange(B, device=dev)
            pos = torch.full((B,), ids.shape[1] - 1, device=dev,
                             dtype=torch.long)
            S = looper.band(X, calls)  # settled prefix + provisional new pos
            S, X = carry_steps(looper, adapter, e, calls, S, X, [(rows, pos)])
        X_store = X
        logits = looper.suffix_logits(S, calls, last_only=True)
        nxt = logits[:, -1].argmax(-1)
        nxt = torch.where(done, torch.full_like(nxt, list(eos)[0]), nxt)
        ids = torch.cat([ids, nxt[:, None]], 1)
        mask = torch.cat([mask, (~done)[:, None].long()], 1)
        done |= torch.tensor([t.item() in eos for t in nxt], device=dev)
        if done.all():
            break
    return ids