"""Block-structured attention masks for Abstract-CoT bottlenecked SFT.

Training sequence  s = [X ; C ; Z ; Y]:
    X  prompt              C  verbal CoT (teacher rationale)
    Z  abstract trace      Y  answer

The paper's bottleneck (Ramji et al., S3.2, eq. around the block mask A) is *exactly*
standard causal masking with ONE edge removed: **answer positions may not attend to the
verbal-CoT positions**. Every bit of C-information reaching Y must route through the
abstract hidden states H_Z, giving the Markov structure  C -> H_Z -> Y  and the channel
bound  I(C;Y|X,Z) <= I(C;H_Z|X,Z)  that scales with the abstract length m.

Note Z attending X u C u Z_<=i is *already* causal (C precedes Z), so it is NOT a
deviation -- the only deviation is the forbidden (Y query, C key) block. We encode roles
per token and forbid a configurable set of (query_role, key_role) pairs, then materialise
a 4D additive mask. transformers>=5.8 eager/sdpa attention honors the 4D mask bit-exactly
(verified: tests/test_masking.py).
"""
from __future__ import annotations

import torch

# per-token roles
X, C, Z, Y = 0, 1, 2, 3
PAD = -1
ROLE = {"X": X, "C": C, "Z": Z, "Y": Y, "PAD": PAD}

# the paper's bottleneck: the answer cannot see the verbal CoT
BOTTLENECK_FORBIDDEN = ((Y, C),)
# TIGHT bottleneck: the answer attends ONLY the abstract -- it cannot see the prompt X OR the
# verbal CoT C. Forces a strict prompt -> abstract -> answer Markov chain, so the abstract is
# load-bearing BY CONSTRUCTION (the answer literally has no other path to the question). Use
# with abstract-only answer generation at inference (no prompt in the answer's context).
TIGHT_FORBIDDEN = ((Y, C), (Y, X))
# debug-only: also forbid Z->C, fully isolating C so Y becomes invariant to C (used to
# prove the 4D mask is actually honored end-to-end)
ISOLATE_C_FORBIDDEN = ((Y, C), (Z, C))


def segment_roles(x_len: int, c_len: int, z_len: int, y_len: int, pad_to: int | None = None) -> torch.Tensor:
    """Build a [L] role vector for one example laid out as [X C Z Y] (+ right PAD)."""
    roles = [X] * x_len + [C] * c_len + [Z] * z_len + [Y] * y_len
    if pad_to is not None:
        assert pad_to >= len(roles), f"pad_to={pad_to} < seq len {len(roles)}"
        roles += [PAD] * (pad_to - len(roles))
    return torch.tensor(roles, dtype=torch.long)


def build_attention_mask(
    role_ids: torch.Tensor,
    dtype: torch.dtype = torch.float32,
    forbidden_pairs=BOTTLENECK_FORBIDDEN,
) -> torch.Tensor:
    """role_ids: LongTensor[B, L] of roles in {X,C,Z,Y,PAD}.

    Returns an additive attention mask FloatTensor[B, 1, L, L] (0 = attend, dtype-min =
    masked), ready to pass to a Qwen3 model as ``attention_mask`` under eager/sdpa attn.
    """
    assert role_ids.dim() == 2, "role_ids must be [B, L]"
    B, L = role_ids.shape
    dev = role_ids.device
    neg = torch.finfo(dtype).min

    idx = torch.arange(L, device=dev)
    allowed = (idx[None, :] <= idx[:, None])[None].expand(B, L, L).clone()  # causal [B,L,L]

    q_role = role_ids[:, :, None]   # [B,L,1]
    k_role = role_ids[:, None, :]   # [B,1,L]
    for qr, kr in forbidden_pairs:
        allowed &= ~((q_role == qr) & (k_role == kr))
    allowed &= (k_role != PAD)      # never attend to padding keys

    add = torch.zeros((B, L, L), dtype=dtype, device=dev)
    add.masked_fill_(~allowed, neg)

    # PAD *queries* would otherwise be all-masked -> NaN softmax; let them attend self.
    pad_q = (role_ids == PAD)
    if pad_q.any():
        eye = torch.eye(L, dtype=torch.bool, device=dev)[None].expand(B, L, L)
        add = torch.where(pad_q[:, :, None] & eye, torch.zeros((), dtype=dtype, device=dev), add)

    return add[:, None]             # [B,1,L,L]