# Copyright 2026 Modilify
# SPDX-License-Identifier: LicenseRef-Modilify-Open-Model-1.0
"""Fixed-shape latent deliberation state for Modilify Mk1 decoding.

The state deliberately contains no vocabulary-sized tensors.  Keeping the
per-canvas information in a small latent space prevents iterative diffusion
rollouts from retaining one logits/probability allocation per denoise pass.
"""

from __future__ import annotations

from dataclasses import dataclass

import torch
from torch import nn


@dataclass
class LatentDeliberationState:
    """Persistent, fixed-size state for one or more canvas episodes."""

    token_latents: torch.Tensor
    memory_slots: torch.Tensor
    confidence: torch.Tensor
    entropy: torch.Tensor
    age: torch.Tensor
    token_changed: torch.Tensor
    confidence_delta: torch.Tensor
    entropy_delta: torch.Tensor
    ponder_steps: torch.Tensor
    stagnation_steps: torch.Tensor

    @classmethod
    def empty(
        cls,
        *,
        batch_size: int,
        canvas_length: int,
        latent_dim: int,
        memory_slots: int,
        device: torch.device,
        dtype: torch.dtype,
    ) -> "LatentDeliberationState":
        """Create a zero-initialized recurrent state.

        Args:
            batch_size: Number of independent sequences.
            canvas_length: Number of rolling canvas positions.
            latent_dim: Width of each latent token and memory slot.
            memory_slots: Number of persistent memory slots.
            device: Allocation device.
            dtype: Floating-point dtype for latent tensors.

        Returns:
            A zero-initialized state with integer progress clocks.
        """

        return cls(
            token_latents=torch.zeros(
                batch_size, canvas_length, latent_dim, device=device, dtype=dtype
            ),
            memory_slots=torch.zeros(
                batch_size, memory_slots, latent_dim, device=device, dtype=dtype
            ),
            confidence=torch.zeros(
                batch_size, canvas_length, device=device, dtype=torch.float32
            ),
            entropy=torch.zeros(
                batch_size, canvas_length, device=device, dtype=torch.float32
            ),
            age=torch.zeros(
                batch_size, canvas_length, device=device, dtype=torch.int32
            ),
            token_changed=torch.zeros(
                batch_size, canvas_length, device=device, dtype=torch.float32
            ),
            confidence_delta=torch.zeros(
                batch_size, canvas_length, device=device, dtype=torch.float32
            ),
            entropy_delta=torch.zeros(
                batch_size, canvas_length, device=device, dtype=torch.float32
            ),
            ponder_steps=torch.zeros(batch_size, device=device, dtype=torch.int32),
            stagnation_steps=torch.zeros(batch_size, device=device, dtype=torch.int32),
        )


def advance_trajectory_clocks(
    ponder_steps: torch.Tensor,
    stagnation_steps: torch.Tensor,
    *,
    commit_lengths: torch.LongTensor,
    active_rows: torch.BoolTensor,
    progress_scores: torch.Tensor,
    min_progress: float,
) -> tuple[torch.IntTensor, torch.IntTensor]:
    """Advance useful-ponder and true-stagnation clocks for each row.

    Args:
        ponder_steps: Total waiting steps for each row.
        stagnation_steps: Consecutive non-improving steps for each row.
        commit_lengths: Number of committed tokens for each row.
        active_rows: Rows that are still generating.
        progress_scores: Signed fused-risk improvements.
        min_progress: Smallest improvement that resets stagnation.

    Returns:
        Updated ponder and stagnation counters.
    """

    if min_progress < 0:
        raise ValueError("`min_progress` must be non-negative.")
    if not (
        ponder_steps.shape == stagnation_steps.shape == commit_lengths.shape
        == active_rows.shape == progress_scores.shape
    ):
        raise ValueError("Trajectory clock inputs must share shape [batch].")
    committed = commit_lengths.gt(0)
    waiting = active_rows & ~committed
    improving = progress_scores.ge(min_progress)
    next_ponder = torch.where(
        committed, torch.zeros_like(ponder_steps), ponder_steps + waiting.to(torch.int32)
    )
    next_stagnation = torch.where(
        committed,
        torch.zeros_like(stagnation_steps),
        torch.where(
            waiting & improving,
            torch.zeros_like(stagnation_steps),
            stagnation_steps + waiting.to(torch.int32),
        ),
    )
    return next_ponder.to(torch.int32), next_stagnation.to(torch.int32)


def should_force_trajectory_jump(
    ponder_steps: torch.Tensor,
    stagnation_steps: torch.Tensor,
    *,
    max_ponder_steps: int,
    stagnation_threshold: int,
) -> torch.BoolTensor:
    """Return rows that exhausted either inference progress clock.

    Args:
        ponder_steps: Total waiting steps for each row.
        stagnation_steps: Consecutive non-improving steps for each row.
        max_ponder_steps: Maximum allowed waiting steps.
        stagnation_threshold: Maximum consecutive stagnation steps.

    Returns:
        Boolean mask selecting rows that must use a forced jump.
    """

    if max_ponder_steps <= 0 or stagnation_threshold <= 0:
        raise ValueError("Trajectory jump limits must be positive.")
    return ponder_steps.ge(max_ponder_steps) | stagnation_steps.ge(stagnation_threshold)


class _TemporalTransformerCell(nn.Module):
    """One-step recurrent token update with fixed-slot memory attention."""

    def __init__(
        self, latent_dim: int, num_heads: int, dropout: float,
        local_attention_window: int,
    ) -> None:
        super().__init__()
        self.state_norm = nn.LayerNorm(latent_dim)
        self.observation_norm = nn.LayerNorm(latent_dim)
        # A slot's learned identity is only used for attention addressing.  The
        # recurrent state itself remains pure memory content so commit shifts
        # cannot accidentally write positional identity into persistent state.
        self.memory_address_norm = nn.LayerNorm(latent_dim)
        self.memory_value_norm = nn.LayerNorm(latent_dim)
        self.temporal_update = nn.Linear(2 * latent_dim, 2 * latent_dim)
        self.local_attention = nn.MultiheadAttention(
            latent_dim, num_heads, dropout=dropout, batch_first=True
        )
        self.local_attention_window = local_attention_window
        self.register_buffer("_local_attention_mask", torch.empty(0), persistent=False)
        self.token_memory_attention = nn.MultiheadAttention(
            latent_dim, num_heads, dropout=dropout, batch_first=True
        )
        self.memory_token_attention = nn.MultiheadAttention(
            latent_dim, num_heads, dropout=dropout, batch_first=True
        )
        self.token_ff_norm = nn.LayerNorm(latent_dim)
        self.memory_ff_norm = nn.LayerNorm(latent_dim)
        self.stored_token_norm = nn.LayerNorm(latent_dim)
        self.stored_memory_norm = nn.LayerNorm(latent_dim)
        expansion = latent_dim * 4
        self.token_ff = nn.Sequential(
            nn.Linear(latent_dim, expansion),
            nn.SiLU(),
            nn.Linear(expansion, latent_dim),
        )
        self.memory_ff = nn.Sequential(
            nn.Linear(latent_dim, expansion),
            nn.SiLU(),
            nn.Linear(expansion, latent_dim),
        )

    def forward(
        self,
        previous_tokens: torch.Tensor,
        observation: torch.Tensor,
        memory: torch.Tensor,
        memory_slot_identity: torch.Tensor,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        gate_logits, candidate = self.temporal_update(
            torch.cat(
                (self.state_norm(previous_tokens), self.observation_norm(observation)),
                dim=-1,
            )
        ).chunk(2, dim=-1)
        gate = torch.sigmoid(gate_logits)
        tokens = gate * previous_tokens + (1.0 - gate) * torch.nn.functional.silu(candidate)
        if (
            self._local_attention_mask.shape != (tokens.shape[1], tokens.shape[1])
            or self._local_attention_mask.device != tokens.device
            or self._local_attention_mask.dtype != tokens.dtype
        ):
            positions = torch.arange(tokens.shape[1], device=tokens.device)
            allowed = (
                positions[:, None] - positions[None, :]
            ).abs() < self.local_attention_window
            self._local_attention_mask = torch.zeros(
                tokens.shape[1], tokens.shape[1], device=tokens.device, dtype=tokens.dtype
            ).masked_fill(~allowed, torch.finfo(tokens.dtype).min)
        local_update, _ = self.local_attention(
            self.state_norm(tokens), self.state_norm(tokens), self.state_norm(tokens),
            attn_mask=self._local_attention_mask, need_weights=False,
        )
        tokens = tokens + local_update

        addressed_memory = self.memory_address_norm(memory + memory_slot_identity)
        memory_values = self.memory_value_norm(memory)
        token_memory_update, _ = self.token_memory_attention(
            self.state_norm(tokens), addressed_memory, memory_values, need_weights=False
        )
        tokens = tokens + token_memory_update
        tokens = tokens + self.token_ff(self.token_ff_norm(tokens))

        memory_token_update, _ = self.memory_token_attention(
            addressed_memory,
            self.state_norm(tokens),
            self.state_norm(tokens),
            need_weights=False,
        )
        memory = memory + memory_token_update
        memory = memory + self.memory_ff(self.memory_ff_norm(memory))
        # This module is a recurrent cell, not a depth-only Transformer block.
        # Persist normalized state so repeated denoise updates cannot accumulate
        # an unbounded residual magnitude across time.
        return self.stored_token_norm(tokens), self.stored_memory_norm(memory)


class LatentDeliberationTransformer(nn.Module):
    """Small recurrent Transformer that compresses repeated denoise context."""

    def __init__(
        self,
        *,
        hidden_size: int,
        latent_dim: int = 512,
        memory_slots: int = 16,
        num_layers: int = 2,
        num_heads: int = 8,
        local_attention_window: int = 32,
        dropout: float = 0.0,
    ) -> None:
        super().__init__()
        if latent_dim % num_heads:
            raise ValueError("`latent_dim` must be divisible by `num_heads`.")
        if local_attention_window <= 0:
            raise ValueError("`local_attention_window` must be positive.")
        self.hidden_size = hidden_size
        self.latent_dim = latent_dim
        self.memory_slots = memory_slots
        self.heavy_projection = nn.Linear(hidden_size, latent_dim, bias=False)
        self.embedding_projection = nn.Linear(hidden_size, latent_dim, bias=False)
        self.scalar_projection = nn.Linear(11, latent_dim, bias=False)
        self.blocks = nn.ModuleList(
            [
                _TemporalTransformerCell(
                    latent_dim, num_heads, dropout, local_attention_window
                )
                for _ in range(num_layers)
            ]
        )
        self.output_norm = nn.LayerNorm(latent_dim)
        self.output_projection = nn.Linear(latent_dim, hidden_size, bias=False)
        self.memory_slot_identity = nn.Parameter(torch.empty(memory_slots, latent_dim))
        self.reset_memory_slot_identity()

    @torch.no_grad()
    def reset_memory_slot_identity(self) -> None:
        """Restore learned memory addresses after generic initialization."""

        nn.init.normal_(self.memory_slot_identity, mean=0.0, std=0.02)

    def project_context(self, token_latents: torch.Tensor) -> torch.Tensor:
        """Translate latent state into a self-conditioning embedding."""

        normalized_tokens = self.output_norm(token_latents)
        return self.output_projection(normalized_tokens)

    def forward(
        self,
        *,
        heavy_hidden: torch.Tensor,
        token_embeddings: torch.Tensor,
        confidence: torch.Tensor,
        entropy: torch.Tensor,
        state: LatentDeliberationState,
    ) -> tuple[torch.Tensor, LatentDeliberationState]:
        """Advance latent memory and produce decoder self-conditioning.

        Args:
            heavy_hidden: Hidden states from the previous decoder pass.
            token_embeddings: Embeddings of current noisy canvas tokens.
            confidence: Proposal confidence for each canvas position.
            entropy: Proposal entropy for each canvas position.
            state: Persistent latent state from the preceding pass.

        Returns:
            Self-conditioning embeddings and the next compact latent state.
        """

        if heavy_hidden.ndim != 3:
            raise ValueError("`heavy_hidden` must have shape [batch, canvas, hidden].")
        if heavy_hidden.shape != token_embeddings.shape:
            raise ValueError("`heavy_hidden` and `token_embeddings` must have the same shape.")
        batch_size, canvas_length, hidden_size = heavy_hidden.shape
        if hidden_size != self.hidden_size:
            raise ValueError("Unexpected hidden size for latent deliberation.")
        expected_state = (batch_size, canvas_length, self.latent_dim)
        if state.token_latents.shape != expected_state:
            raise ValueError("State token latents do not match the current canvas.")
        if state.memory_slots.shape != (batch_size, self.memory_slots, self.latent_dim):
            raise ValueError("State memory slots do not match this module.")
        if state.age.dtype is not torch.int32:
            raise TypeError("Latent deliberation ages must use int32.")

        scalars = torch.stack(
            (
                confidence.to(dtype=heavy_hidden.dtype),
                entropy.to(dtype=heavy_hidden.dtype).log1p(),
                state.age.to(dtype=heavy_hidden.dtype).clamp_max(32767).log1p(),
                torch.linspace(
                    -1.0, 1.0, canvas_length, device=heavy_hidden.device,
                    dtype=heavy_hidden.dtype,
                ).unsqueeze(0).expand(batch_size, -1),
                state.token_changed.to(dtype=heavy_hidden.dtype),
                state.confidence_delta.to(dtype=heavy_hidden.dtype),
                state.entropy_delta.to(dtype=heavy_hidden.dtype).sign()
                * state.entropy_delta.to(dtype=heavy_hidden.dtype).abs().log1p(),
                state.ponder_steps.to(dtype=heavy_hidden.dtype).log1p()[:, None]
                .expand(-1, canvas_length),
                state.stagnation_steps.to(dtype=heavy_hidden.dtype).log1p()[:, None]
                .expand(-1, canvas_length),
                confidence.to(dtype=heavy_hidden.dtype)
                * torch.exp(-entropy.to(dtype=heavy_hidden.dtype).clamp_min(0.0)),
                state.confidence_delta.to(dtype=heavy_hidden.dtype).clamp_min(0.0)
                + (-state.entropy_delta.to(dtype=heavy_hidden.dtype)).clamp_min(0.0).log1p(),
            ),
            dim=-1,
        )
        observation = (
            self.heavy_projection(heavy_hidden)
            + self.embedding_projection(token_embeddings)
            + self.scalar_projection(scalars)
        )
        tokens = state.token_latents
        memory = state.memory_slots
        slot_identity = self.memory_slot_identity.to(device=memory.device, dtype=memory.dtype)
        slot_identity = slot_identity.unsqueeze(0).expand(batch_size, -1, -1)
        for block in self.blocks:
            tokens, memory = block(tokens, observation, memory, slot_identity)
            observation = tokens

        context = self.project_context(tokens)
        next_state = LatentDeliberationState(
            token_latents=tokens,
            memory_slots=memory,
            confidence=confidence.to(dtype=torch.float32),
            entropy=entropy.to(dtype=torch.float32),
            age=state.age,
            token_changed=state.token_changed,
            confidence_delta=state.confidence_delta,
            entropy_delta=state.entropy_delta,
            ponder_steps=state.ponder_steps,
            stagnation_steps=state.stagnation_steps,
        )
        return context, next_state


__all__ = [
    "LatentDeliberationState", "LatentDeliberationTransformer",
    "advance_trajectory_clocks", "should_force_trajectory_jump",
]