"""Soft-token (continuous-thought) machinery for latent_threads — CODI-style feedback.

The latent block is L SOFT positions: the input embedding of z_{t+1} is a learned projection of
the LAST-LAYER hidden state at z_t (z_1's input = proj(h at the last prefix token)). Unlike the
dot organisms (constant input token, information moves between latent positions only via
attention), this is a genuine recurrence through the embedding channel — each z_t is a vector
the model wrote, and the hypothesis is that z_t holds the PARALLEL states of all threads after
step t. Same Y!->X bottleneck mask (query+answer never see the prompt) and delayed query.

Implementation: iterative full-sequence forwards with inputs_embeds (L+1 passes). Pass k fixes
the embedding of z_{k}; all passes stay in-graph, so the answer CE backpropagates through the
whole latent chain (BPTT). Sequences are ~120-250 tokens, so L+1 eager forwards are cheap.
"""
from __future__ import annotations

import torch
import torch.nn as nn

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

ORGANISM = ((ROLE_Y, ROLE_X),)                     # answer reads only the latent block
ABLATED = ((ROLE_Y, ROLE_X), (ROLE_Z, ROLE_X))     # + latents blinded to prompt (control => chance)
# TIGHT recurrence: latents ALSO cannot attend the prompt, so the prompt reaches the computation
# ONLY through z_1's projected input embedding -> the soft-token RECURRENCE is load-bearing by
# construction. (The default ORGANISM lets latents recompute from the prompt via Z->X, which the
# {zero-all,random,cross-patch} controls revealed makes the input feedback vestigial.)
TIGHT = ((ROLE_Y, ROLE_X), (ROLE_Z, ROLE_X))


class Projection(nn.Module):
    """h_lastlayer -> next input embedding. LN front, small-init output => soft inputs start ~0."""

    def __init__(self, d: int, dtype=torch.bfloat16):
        super().__init__()
        self.net = nn.Sequential(nn.LayerNorm(d), nn.Linear(d, d), nn.GELU(), nn.Linear(d, d))
        nn.init.normal_(self.net[-1].weight, std=1e-3)
        nn.init.zeros_(self.net[-1].bias)
        self.net.to(dtype)

    def forward(self, h):
        return self.net(h)


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


def build_soft_mask(roles, mode, dtype=torch.bfloat16):
    """mode: 'organism' (Y!->X), 'tight' (Y!->X & Z!->X), 'tight_first' (tight, but the FIRST
    latent position MAY attend the prompt — read the problem once, then recur in latents only =
    genuine load-bearing recurrence on lookup tasks). Position-aware, so it post-edits the mask."""
    if mode == "organism":
        return build_attention_mask(roles, dtype=dtype, forbidden_pairs=ORGANISM)
    add = build_attention_mask(roles, dtype=dtype, forbidden_pairs=TIGHT)  # [B,1,L,L]
    if mode == "tight":
        return add
    if mode != "tight_first":
        raise ValueError(mode)
    B, _, L, _ = add.shape
    for b in range(B):
        zpos = (roles[b] == ROLE_Z).nonzero(as_tuple=True)[0]
        if len(zpos) == 0:
            continue
        z0 = int(zpos[0])
        xkeys = (roles[b] == ROLE_X).nonzero(as_tuple=True)[0]
        causal = xkeys[xkeys <= z0]
        add[b, 0, z0, causal] = 0.0  # re-allow z_1 -> (causal) prompt
    return add


def build_soft_batch(tok, task, probs, L, im_end, device, with_answer=True):
    """Padded batch pieces for the soft layout [X: prefix+<think>\\n][Z: L soft][Y: close+query
    (+answer)]. Returns (input_ids with z slots = pad, roles, z_starts, label_starts, lengths)."""
    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 + ([T.digit_ids(tok)[task.answer(p)]] + _ids(tok, "}") + [im_end] if with_answer else [])
        ids = x + [tok.pad_token_id] * L + y
        roles = [ROLE_X] * len(x) + [ROLE_Z] * L + [ROLE_Y] * len(y)
        rows.append((ids, roles, len(x), len(x) + L + 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 _mask(roles, forbidden, dtype=torch.bfloat16):
    """forbidden may be a tuple of (q,k) pairs (ORGANISM/ABLATED/TIGHT) OR a mode string
    ('organism'/'tight'/'tight_first') for position-aware masks."""
    if isinstance(forbidden, str):
        return build_soft_mask(roles, forbidden, dtype=dtype)
    return build_attention_mask(roles, dtype=dtype, forbidden_pairs=forbidden)

def soft_forward(model, proj, tok, task, probs, L, im_end, device, with_answer=True,
                 forbidden=ORGANISM, collect_z=False, collect_fb=False):
    """L+1 in-graph forwards filling the soft slots; returns (logits, labels_meta, z_states[, fb]).

    z_states (if collect_z): [B, L, D] last-layer hidden at the soft positions (the vectors fed
    back) — the organism's latent trace. fb (if collect_fb): [B, L, D] the PROJECTED feedback
    vectors proj(h(z_t)) for t=1..L (what the model "wrote" after each latent step; the last one
    is computed but never consumed) — the supervision site for ground-truth latent thoughts."""
    input_ids, roles, z_starts, label_starts, lengths = build_soft_batch(
        tok, task, probs, L, im_end, device, with_answer)
    B, Lmax = input_ids.shape
    attn4d = _mask(roles, forbidden)
    pos = torch.arange(Lmax, device=device)[None].expand(B, Lmax)
    emb_layer = model.get_input_embeddings()
    E = emb_layer(input_ids)  # [B, Lmax, D]; z slots hold pad embeddings until filled
    bidx = torch.arange(B, device=device)
    zs = torch.tensor(z_starts, device=device)
    out, fbs = None, []
    for t in range(L + 1):
        out = model(inputs_embeds=E, attention_mask=attn4d, position_ids=pos,
                    output_hidden_states=True)
        if t == L:
            break
        h = out.hidden_states[-1][bidx, zs + t - 1] if t > 0 else out.hidden_states[-1][bidx, zs - 1]
        z_in = proj(h)
        if t > 0:
            fbs.append(z_in)  # proj(h(z_t)) for t=1..L-1
        E = E.clone()
        E[bidx, zs + t] = z_in
    fb = None
    if collect_fb:
        fbs.append(proj(out.hidden_states[-1][bidx, zs + L - 1]))  # proj(h(z_L)), unconsumed
        fb = torch.stack(fbs, dim=1)  # [B, L, D]
    z_states = None
    if collect_z:
        # "last" = the vectors actually fed back (the latent trace); "mid" = layer-27 residuals
        # at the soft positions (the AO read layer; what AVBench context_activations carry).
        z_states = {
            "last": torch.stack([out.hidden_states[-1][bidx, zs + t] for t in range(L)], dim=1),
            "mid": torch.stack([out.hidden_states[27][bidx, zs + t] for t in range(L)], dim=1),
        }
    if collect_fb:
        return out.logits, (z_starts, label_starts, lengths, input_ids), z_states, fb
    return out.logits, (z_starts, label_starts, lengths, input_ids), z_states


def make_binding(d: int, n_threads: int, device, seed: int = 7):
    """Fixed random sign/permutation operators R_b (orthogonal, O(D) memory) for thread binding."""
    g = torch.Generator(device="cpu").manual_seed(seed)
    perms = [torch.randperm(d, generator=g).to(device) for _ in range(n_threads)]
    signs = [(torch.randint(0, 2, (d,), generator=g) * 2 - 1).to(device) for _ in range(n_threads)]
    return perms, signs


def gt_latents(task, probs, emb_weight, did, perms, signs, device):
    """Designed ground-truth latent thoughts: z*_t = sum_b R_b E[digit s_b(t)] / sqrt(B) — a
    vocabulary-space superposition of all threads' step-t states with thread binding
    (Latent-SFT-style targets, exact because the generators define the states)."""
    import math
    states = torch.tensor([task.step_states(p) for p in probs], device=device)  # [B, L, n_threads]
    did_t = torch.tensor(did, device=device)
    Bn = states.shape[-1]
    tgt = torch.zeros(states.shape[0], states.shape[1], emb_weight.shape[1],
                      dtype=torch.float32, device=device)
    for b in range(Bn):
        e = emb_weight[did_t[states[:, :, b]]].float()    # [B, L, D]
        tgt += signs[b].float() * e[..., perms[b]]
    return tgt / math.sqrt(Bn)


@torch.no_grad()
def soft_latent_inputs(model, proj, tok, task, probs, L, im_end, device, forbidden=ORGANISM):
    """Run the recurrence and return the FILLED input embeddings + metadata, so controls can
    manipulate the soft-token inputs (shuffle / cross-patch / leave-out) and re-read the answer.
    Returns (E, attn4d, pos, z_starts, lengths)."""
    input_ids, roles, z_starts, label_starts, lengths = build_soft_batch(
        tok, task, probs, L, im_end, device, with_answer=False)
    B, Lmax = input_ids.shape
    attn4d = _mask(roles, forbidden)
    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)
    for t in range(L):
        out = model(inputs_embeds=E, attention_mask=attn4d, position_ids=pos, output_hidden_states=True)
        h = out.hidden_states[-1][bidx, zs + t - 1] if t > 0 else out.hidden_states[-1][bidx, zs - 1]
        E = E.clone()
        E[bidx, zs + t] = proj(h)
    return E, attn4d, pos, z_starts, lengths


@torch.no_grad()
def _read_from_E(model, tok, E, attn4d, pos, lengths, device):
    did = torch.tensor(T.digit_ids(tok), device=device)
    logits = model(inputs_embeds=E, attention_mask=attn4d, position_ids=pos).logits
    last = logits[torch.arange(E.shape[0], device=device), torch.tensor(lengths, device=device) - 1]
    return did[last[:, did].argmax(-1)].tolist()


@torch.no_grad()
def soft_controls(model, proj, tok, task, probs, L, im_end, device, seed: int = 0):
    """{shuffle, cross-patch, leave-out} CoT controls on the soft latent block."""
    did = T.digit_ids(tok)
    id2d = {t: d for d, t in enumerate(did)}
    E, attn4d, pos, z_starts, lengths = soft_latent_inputs(model, proj, tok, task, probs, L, im_end, device)
    bidx = torch.arange(E.shape[0], device=device)
    zs = torch.tensor(z_starts, device=device)
    gold = [task.answer(p) for p in probs]

    def acc(pred, ref):
        return sum(int(id2d.get(pr, -1) == r) for pr, r in zip(pred, ref)) / len(ref)

    base = acc(_read_from_E(model, tok, E, attn4d, pos, lengths, device), gold)

    # shuffle: permute the L soft-token embeddings within each example (deterministic per-row)
    Esh = E.clone()
    g = torch.Generator(device="cpu").manual_seed(seed)
    for j in range(E.shape[0]):
        perm = torch.randperm(L, generator=g)
        Esh[j, z_starts[j]:z_starts[j] + L] = E[j, z_starts[j]:z_starts[j] + L][perm]
    shuf = acc(_read_from_E(model, tok, Esh, attn4d, pos, lengths, device), gold)

    # cross-patch: splice a DIFFERENT instance's soft block into each example (donor = roll by 1).
    # Under Y!->X the answer sees only Z -> it should follow the DONOR's state for this query slot.
    Ecp = E.clone()
    donor = [(j + 1) % E.shape[0] for j in range(E.shape[0])]
    for j in range(E.shape[0]):
        Ecp[j, z_starts[j]:z_starts[j] + L] = E[donor[j], z_starts[donor[j]]:z_starts[donor[j]] + L]
    cp_pred = _read_from_E(model, tok, Ecp, attn4d, pos, lengths, device)
    # donor-following GT = the donor instance answered with THIS example's query
    from dataclasses import replace as _rep
    donor_gold = [task.answer(_rep(probs[donor[j]], q=probs[j].q)) for j in range(len(probs))]
    cp_donor = acc(cp_pred, donor_gold)
    cp_orig = acc(cp_pred, gold)

    # leave-out: zero each soft position in turn; report mean acc and worst position
    lo = []
    for k in range(L):
        Elo = E.clone()
        Elo[bidx, zs + k] = 0.0
        lo.append(acc(_read_from_E(model, tok, Elo, attn4d, pos, lengths, device), gold))

    # positive checks: zero ALL soft inputs / replace with random. If acc stays high, the soft-token
    # INPUT recurrence is vestigial (the latent positions recompute from the prompt via Z->X attn).
    Ez = E.clone()
    for j in range(E.shape[0]):
        Ez[j, z_starts[j]:z_starts[j] + L] = 0.0
    zero_all = acc(_read_from_E(model, tok, Ez, attn4d, pos, lengths, device), gold)
    Er = E.clone()
    sd = float(E.float().std())
    for j in range(E.shape[0]):
        noise = (torch.randn(L, E.shape[-1], generator=g, dtype=torch.float32) * sd).to(device, E.dtype)
        Er[j, z_starts[j]:z_starts[j] + L] = noise
    rand_all = acc(_read_from_E(model, tok, Er, attn4d, pos, lengths, device), gold)

    return {"baseline": base, "shuffle": shuf, "crosspatch_donor": cp_donor,
            "crosspatch_orig": cp_orig, "leaveout_mean": sum(lo) / L,
            "leaveout_min": min(lo), "leaveout_per_pos": lo,
            "zero_all_inputs": zero_all, "random_inputs": rand_all}


@torch.no_grad()
def soft_readout_acc(model, proj, tok, task, probs, L, im_end, device, forbidden=ORGANISM):
    model.eval()
    did = torch.tensor(T.digit_ids(tok), device=device)
    correct, bs = 0, 16
    for i in range(0, len(probs), bs):
        batch = probs[i : i + bs]
        logits, (_, _, lengths, _), _ = soft_forward(model, proj, tok, task, batch, L, im_end,
                                                     device, with_answer=False, forbidden=forbidden)
        last = logits[torch.arange(len(batch), device=device), torch.tensor(lengths, device=device) - 1]
        pred = last[:, did].argmax(dim=-1).tolist()
        for j, p in enumerate(batch):
            correct += int(pred[j] == task.answer(p))
    model.train()
    return correct / len(probs)


@torch.no_grad()
def soft_completeness(model, proj, tok, task, probs, L, im_end, device):
    from latent_threads.eval_masked import _clone_with_query
    did = torch.tensor(T.digit_ids(tok), device=device)
    n_q = len(task.all_queries(probs[0]))
    correct = [0] * n_q
    for p in probs:
        variants = [(_clone_with_query(task, p, qi), ans) for qi, (_, ans) in enumerate(task.all_queries(p))]
        logits, (_, _, lengths, _), _ = soft_forward(model, proj, tok, task, [v for v, _ in variants],
                                                     L, im_end, device, with_answer=False)
        last = logits[torch.arange(len(variants), device=device), torch.tensor(lengths, device=device) - 1]
        pred = last[:, did].argmax(dim=-1).tolist()
        for qi, (_, ans) in enumerate(variants):
            correct[qi] += int(pred[qi] == ans)
    return [c / len(probs) for c in correct]