"""Attention layouts for non-recurrent visual latent reasoning chains.

The sequence is laid out as::

    [multimodal prompt ; clean question ; LOOK_1 ; THINK_1 ; ... ; answer]

All rows are processed once by the native VLM decoder.  A block-sparse causal
graph makes LOOK rows the only latent rows with access to the multimodal
prefix, while THINK rows integrate one visual read without directly seeing
that prefix.  Answer rows can see the clean question and THINK rows, but never
the multimodal prefix or LOOK rows.  Consequently the image-dependent answer
path is

    image -> LOOK_k -> THINK_k -> answer,

without a tied recurrent cell, a visual-state update, or an aggregation
module.

The transition-conditioned variant tightens the graph without adding a
module.  The first LOOK/THINK pair bootstraps from the complete prompt and the
clean question.  Every later pair sees the causal latent prefix, while its
LOOK row receives only visual placeholder rows from the multimodal prefix.
Thus later pairs cannot independently re-solve the original image/question
prompt; their task semantics must arrive through earlier latent states.  The
last THINK is the sole latent exposed to answer rows and therefore acts as the
native-transformer aggregation state.
"""

from __future__ import annotations

from dataclasses import dataclass

import torch


@dataclass(frozen=True)
class PerceiveDeliberateLayout:
    """Physical row layout of one padded training/generation sequence."""

    multimodal_length: int
    question_length: int
    num_pairs: int
    answer_input_length: int = 0

    def __post_init__(self) -> None:
        if self.multimodal_length < 1:
            raise ValueError("multimodal_length must be positive")
        if self.question_length < 1:
            raise ValueError("question_length must be positive")
        if self.num_pairs < 1:
            raise ValueError("num_pairs must be positive")
        if self.answer_input_length < 0:
            raise ValueError("answer_input_length must be non-negative")

    @property
    def num_latent_tokens(self) -> int:
        return 2 * self.num_pairs

    @property
    def question_start(self) -> int:
        return self.multimodal_length

    @property
    def latent_start(self) -> int:
        return self.multimodal_length + self.question_length

    @property
    def answer_start(self) -> int:
        return self.latent_start + self.num_latent_tokens

    @property
    def sequence_length(self) -> int:
        return self.answer_start + self.answer_input_length

    @property
    def question_slice(self) -> slice:
        return slice(self.question_start, self.latent_start)

    @property
    def answer_slice(self) -> slice:
        return slice(self.answer_start, self.sequence_length)

    @property
    def look_indices(self) -> tuple[int, ...]:
        return tuple(self.latent_start + 2 * index for index in range(self.num_pairs))

    @property
    def think_indices(self) -> tuple[int, ...]:
        return tuple(index + 1 for index in self.look_indices)


def build_perceive_deliberate_mask(
    layout: PerceiveDeliberateLayout,
    multimodal_attention_mask: torch.Tensor,
    question_attention_mask: torch.Tensor,
    answer_input_attention_mask: torch.Tensor | None = None,
    *,
    multimodal_visual_mask: torch.Tensor | None = None,
    local_visual_chain: bool = False,
    transition_conditioned: bool = False,
    question_visible_through_pair: int = 1,
    disable_chain_links: bool = False,
    disable_late_visual: bool = False,
    final_pair_only: bool = False,
) -> torch.BoolTensor:
    """Build the strict LOOK/THINK/answer visibility graph.

    The returned boolean mask has shape ``[B, 1, L, L]`` and follows PyTorch
    SDPA semantics: ``True`` means that the query may attend to the key.

    In the original graph, LOOK_k sees the valid multimodal prompt, the clean
    question, and (for k>1) only THINK_{k-1}; THINK_k sees the clean question
    and LOOK_k.  In the transition-conditioned graph, pair 1 keeps those
    bootstrap inputs, but every later latent sees its complete causal latent
    prefix, later LOOK rows see only visual placeholder keys, and later THINK
    rows normally receive no independent question shortcut.  For a controlled
    curriculum/intervention, ``question_visible_through_pair`` may temporarily
    retain the clean-question edge through a later pair; ``1`` is the strict
    inference graph.  Answer rows see only the final THINK in that variant.
    Every row also sees itself so padded query
    rows remain numerically defined; padded rows are never exposed as keys to
    semantic rows.
    """

    if multimodal_attention_mask.ndim != 2:
        raise ValueError("multimodal_attention_mask must have shape [B, L_mm]")
    if question_attention_mask.ndim != 2:
        raise ValueError("question_attention_mask must have shape [B, L_q]")
    batch_size = multimodal_attention_mask.shape[0]
    expected_mm = (batch_size, layout.multimodal_length)
    expected_q = (batch_size, layout.question_length)
    if tuple(multimodal_attention_mask.shape) != expected_mm:
        raise ValueError(
            "multimodal mask/layout mismatch: "
            f"expected {expected_mm}, got {tuple(multimodal_attention_mask.shape)}"
        )
    if tuple(question_attention_mask.shape) != expected_q:
        raise ValueError(
            "question mask/layout mismatch: "
            f"expected {expected_q}, got {tuple(question_attention_mask.shape)}"
        )
    if layout.answer_input_length:
        if answer_input_attention_mask is None:
            raise ValueError("answer_input_attention_mask is required")
        expected_answer = (batch_size, layout.answer_input_length)
        if tuple(answer_input_attention_mask.shape) != expected_answer:
            raise ValueError(
                "answer mask/layout mismatch: "
                f"expected {expected_answer}, got "
                f"{tuple(answer_input_attention_mask.shape)}"
            )
    elif answer_input_attention_mask is not None and answer_input_attention_mask.numel():
        raise ValueError("received answer mask for an empty answer-input segment")

    device = multimodal_attention_mask.device
    length = layout.sequence_length
    visible = torch.zeros(
        batch_size, length, length, dtype=torch.bool, device=device
    )
    mm_valid = multimodal_attention_mask.bool()
    q_valid = question_attention_mask.bool()
    visual_valid = None
    if transition_conditioned or local_visual_chain:
        if not 1 <= question_visible_through_pair <= layout.num_pairs:
            raise ValueError(
                "question_visible_through_pair must lie in "
                f"[1,{layout.num_pairs}]"
            )
        if multimodal_visual_mask is None:
            raise ValueError(
                "multimodal_visual_mask is required for the "
                "transition-conditioned graph"
            )
        if tuple(multimodal_visual_mask.shape) != expected_mm:
            raise ValueError(
                "visual mask/layout mismatch: "
                f"expected {expected_mm}, got "
                f"{tuple(multimodal_visual_mask.shape)}"
            )
        visual_valid = multimodal_visual_mask.bool()
        if bool((visual_valid & ~mm_valid).any()):
            raise ValueError("visual keys must be a subset of valid multimodal keys")

    # The source multimodal prompt keeps the native causal graph.  Padded query
    # rows are harmless and receive a self edge below.
    mm_causal = torch.ones(
        layout.multimodal_length,
        layout.multimodal_length,
        dtype=torch.bool,
        device=device,
    ).tril()
    visible[:, : layout.multimodal_length, : layout.multimodal_length] = (
        mm_causal.unsqueeze(0) & mm_valid[:, None, :]
    )

    # The duplicated question is deliberately text-only: it is causally
    # connected only to preceding valid rows of its own segment.
    q_causal = torch.ones(
        layout.question_length,
        layout.question_length,
        dtype=torch.bool,
        device=device,
    ).tril()
    visible[
        :, layout.question_slice, layout.question_slice
    ] = q_causal.unsqueeze(0) & q_valid[:, None, :]

    if local_visual_chain:
        # Homogeneous one-pass latent chain.  Visual memory is persistent and
        # available at every position, while task state has exactly one local
        # predecessor edge.  No latent may consume the complete causal prefix.
        latent_indices = range(layout.latent_start, layout.answer_start)
        for latent_offset, latent in enumerate(latent_indices):
            visible[:, latent, : layout.multimodal_length] = visual_valid
            if latent_offset == 0:
                visible[:, latent, layout.question_slice] = q_valid
            else:
                visible[:, latent, latent - 1] = True

        if layout.answer_input_length:
            answer_valid = answer_input_attention_mask.bool()
            answer_causal = torch.ones(
                layout.answer_input_length,
                layout.answer_input_length,
                dtype=torch.bool,
                device=device,
            ).tril()
            visible[:, layout.answer_slice, layout.question_slice] = q_valid[:, None, :]
            visible[:, layout.answer_slice, layout.answer_start - 1] = True
            visible[:, layout.answer_slice, layout.answer_slice] = (
                answer_causal.unsqueeze(0) & answer_valid[:, None, :]
            )

        diagonal = torch.arange(length, device=device)
        visible[:, diagonal, diagonal] = True
        return visible.unsqueeze(1)

    for pair_index, (look, think) in enumerate(
        zip(layout.look_indices, layout.think_indices)
    ):
        if final_pair_only and pair_index != layout.num_pairs - 1:
            continue

        if transition_conditioned and not final_pair_only:
            # Pair 1 bootstraps the causal latent prefix.  Later pairs cannot
            # recover the task from the original prompt: only image placeholder
            # rows remain visible outside the latent prefix.
            if pair_index == 0:
                visible[:, look, : layout.multimodal_length] = mm_valid
                visible[:, look, layout.question_slice] = q_valid
                visible[:, think, layout.question_slice] = q_valid
            else:
                if not disable_late_visual:
                    visible[:, look, : layout.multimodal_length] = visual_valid
                if not disable_chain_links:
                    visible[:, look, layout.latent_start:look] = True
                    visible[:, think, layout.latent_start:look] = True
                if pair_index < question_visible_through_pair:
                    visible[:, look, layout.question_slice] = q_valid
                    visible[:, think, layout.question_slice] = q_valid

            # The local LOOK -> THINK edge is never an inter-pair chain-link
            # intervention and therefore remains present in every arm.
            visible[:, think, look] = True
        else:
            # Original graph, also used by the final-pair-only capacity
            # control: each active pair may independently consume prompt + Q.
            if not disable_late_visual or pair_index == 0 or final_pair_only:
                visible[:, look, : layout.multimodal_length] = mm_valid
            visible[:, look, layout.question_slice] = q_valid
            if pair_index and not disable_chain_links and not final_pair_only:
                visible[:, look, layout.think_indices[pair_index - 1]] = True

            # Deliberation: no original multimodal key can be consumed here.
            visible[:, think, layout.question_slice] = q_valid
            visible[:, think, look] = True

    if layout.answer_input_length:
        answer_valid = answer_input_attention_mask.bool()
        answer_causal = torch.ones(
            layout.answer_input_length,
            layout.answer_input_length,
            dtype=torch.bool,
            device=device,
        ).tril()
        visible[:, layout.answer_slice, layout.question_slice] = q_valid[:, None, :]
        answer_thinks = (
            [layout.think_indices[-1]]
            if final_pair_only or transition_conditioned
            else list(layout.think_indices)
        )
        visible[:, layout.answer_slice, answer_thinks] = True
        visible[:, layout.answer_slice, layout.answer_slice] = (
            answer_causal.unsqueeze(0) & answer_valid[:, None, :]
        )

    # Avoid all-masked softmax rows for physical padding.  Semantic queries do
    # not receive padded keys because every segment assignment above uses its
    # validity mask.
    diagonal = torch.arange(length, device=device)
    visible[:, diagonal, diagonal] = True
    return visible.unsqueeze(1)


def build_perceive_deliberate_positions(
    layout: PerceiveDeliberateLayout,
    multimodal_position_ids: torch.LongTensor,
    multimodal_attention_mask: torch.Tensor,
    question_attention_mask: torch.Tensor,
    answer_input_attention_mask: torch.Tensor | None = None,
) -> torch.LongTensor:
    """Extend native multimodal M-RoPE with logical text positions.

    The clean question keeps ordinary relative text positions but is shifted
    after the largest valid multimodal M-RoPE coordinate.  Latent and answer
    rows then continue from each item's *logical* question length, independent
    of right padding.  All three M-RoPE coordinates are equal for these
    non-spatial rows.
    """

    if multimodal_position_ids.ndim != 3 or multimodal_position_ids.shape[0] != 3:
        raise ValueError("multimodal_position_ids must have shape [3, B, L_mm]")
    batch_size = multimodal_attention_mask.shape[0]
    if tuple(multimodal_position_ids.shape[1:]) != (
        batch_size,
        layout.multimodal_length,
    ):
        raise ValueError("multimodal position/layout mismatch")

    mm_valid = multimodal_attention_mask.bool()
    masked_mm = multimodal_position_ids.masked_fill(~mm_valid.unsqueeze(0), -1)
    continuation = masked_mm.amax(dim=(0, 2)).clamp_min(-1) + 1  # [B]

    q_valid = question_attention_mask.bool()
    q_relative = q_valid.long().cumsum(dim=-1) - 1
    q_relative = q_relative.masked_fill(~q_valid, 0)
    q_positions = continuation[:, None] + q_relative
    q_lengths = q_valid.sum(dim=-1)

    latent_relative = torch.arange(
        layout.num_latent_tokens,
        device=multimodal_position_ids.device,
    )
    latent_positions = (
        continuation[:, None] + q_lengths[:, None] + latent_relative[None, :]
    )

    segments = [multimodal_position_ids, q_positions.unsqueeze(0).expand(3, -1, -1)]
    segments.append(latent_positions.unsqueeze(0).expand(3, -1, -1))
    if layout.answer_input_length:
        if answer_input_attention_mask is None:
            raise ValueError("answer_input_attention_mask is required")
        answer_relative = torch.arange(
            layout.answer_input_length,
            device=multimodal_position_ids.device,
        )
        answer_positions = (
            continuation[:, None]
            + q_lengths[:, None]
            + layout.num_latent_tokens
            + answer_relative[None, :]
        )
        segments.append(answer_positions.unsqueeze(0).expand(3, -1, -1))
    return torch.cat(segments, dim=-1)


def build_answer_decode_mask(
    layout: PerceiveDeliberateLayout,
    question_attention_mask: torch.Tensor,
    generated_attention_mask: torch.Tensor,
    *,
    local_visual_chain: bool = False,
    final_pair_only: bool = False,
    transition_conditioned: bool = False,
) -> torch.BoolTensor:
    """Visibility for one cached answer query after the latent prefill.

    ``generated_attention_mask`` includes the current answer token.  Physical
    multimodal and LOOK cache rows remain present but are invisible.
    """

    if generated_attention_mask.ndim != 2:
        raise ValueError("generated_attention_mask must have shape [B, N]")
    batch_size, generated_length = generated_attention_mask.shape
    if tuple(question_attention_mask.shape) != (
        batch_size,
        layout.question_length,
    ):
        raise ValueError("question mask/layout mismatch")
    total_keys = layout.answer_start + generated_length
    visible = torch.zeros(
        batch_size, 1, 1, total_keys, dtype=torch.bool,
        device=question_attention_mask.device,
    )
    visible[:, 0, 0, layout.question_slice] = question_attention_mask.bool()
    if local_visual_chain:
        visible[:, 0, 0, layout.answer_start - 1] = True
    else:
        answer_thinks = (
            [layout.think_indices[-1]]
            if final_pair_only or transition_conditioned
            else list(layout.think_indices)
        )
        visible[:, 0, 0, answer_thinks] = True
    visible[:, 0, 0, layout.answer_start:] = generated_attention_mask.bool()
    return visible


__all__ = [
    "PerceiveDeliberateLayout",
    "build_answer_decode_mask",
    "build_perceive_deliberate_mask",
    "build_perceive_deliberate_positions",
]