"""Direct MiniMax H3 DiT block over the standalone Sage3 attention unit."""

import torch
from torch import nn

from .attention import H3SageAttention, rms_norm
from .checkpoint import H3Checkpoint
from .nvfp4 import Nvfp4Linear


def modulate_segments(
    x: torch.Tensor,
    shift: torch.Tensor,
    scale: torch.Tensor,
    segments: list[tuple[int, int, int]],
) -> torch.Tensor:
    """Apply H3's per-modality/per-timestep AdaLN parameters to contiguous rows."""
    output = torch.empty_like(x)
    for start, stop, row in segments:
        output[start:stop] = x[start:stop] * (1 + scale[row].to(x.dtype)) + shift[row].to(x.dtype)
    return output


def gate_segments(
    residual: torch.Tensor,
    update: torch.Tensor,
    gate: torch.Tensor,
    segments: list[tuple[int, int, int]],
) -> torch.Tensor:
    """Add a gated fresh sublayer output without altering the source residual."""
    output = residual.clone()
    for start, stop, row in segments:
        output[start:stop].addcmul_(update[start:stop], gate[row].to(update.dtype))
    return output


class H3SwiGLU(nn.Module):
    def __init__(self, fc1: Nvfp4Linear, fc2: Nvfp4Linear):
        super().__init__()
        self.fc1 = fc1
        self.fc2 = fc2
        self.chunks = 1
        self.chunk_threshold = 4096

    @classmethod
    def from_checkpoint(cls, checkpoint: H3Checkpoint, prefix: str, *, output_dtype=torch.bfloat16):
        return cls(
            checkpoint.nvfp4_linear(f"{prefix}.fc1", output_dtype=output_dtype),
            checkpoint.nvfp4_linear(f"{prefix}.fc2", output_dtype=output_dtype),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if self.chunks > 1 and x.shape[0] >= self.chunk_threshold:
            return torch.cat([self._forward_chunk(chunk) for chunk in x.chunk(self.chunks, dim=0)], dim=0)
        return self._forward_chunk(x)

    def _forward_chunk(self, x: torch.Tensor) -> torch.Tensor:
        gate, up = self.fc1(x).chunk(2, dim=-1)
        return self.fc2(torch.nn.functional.silu(gate).mul_(up))


def configure_mlp_chunking(model: nn.Module, chunks: int, threshold: int = 4096) -> None:
    """Configure exact row-chunked H3 SwiGLU execution to reduce peak activation memory."""
    if chunks < 1:
        raise ValueError("MLP chunks must be >= 1")
    for module in model.modules():
        if isinstance(module, H3SwiGLU):
            module.chunks = chunks
            module.chunk_threshold = threshold


class H3DiTBlock(nn.Module):
    """One H3 transformer block with externally supplied AdaLN tensors."""

    def __init__(
        self,
        norm1_weight: torch.Tensor,
        norm2_weight: torch.Tensor,
        attention: H3SageAttention,
        mlp: H3SwiGLU,
        *,
        norm_eps: float = 1e-5,
    ):
        super().__init__()
        self.attention = attention
        self.mlp = mlp
        self.norm_eps = norm_eps
        self.register_buffer("norm1_weight", norm1_weight, persistent=False)
        self.register_buffer("norm2_weight", norm2_weight, persistent=False)

    @classmethod
    def from_checkpoint(cls, checkpoint: H3Checkpoint, index: int, *, output_dtype=torch.bfloat16, attention_backend: str = "sage2"):
        prefix = f"blocks.{index}"
        return cls(
            checkpoint.tensor(f"{prefix}.norm1.weight", dtype=output_dtype),
            checkpoint.tensor(f"{prefix}.norm2.weight", dtype=output_dtype),
            H3SageAttention.from_checkpoint(checkpoint, f"{prefix}.attn", output_dtype=output_dtype, backend=attention_backend),
            H3SwiGLU.from_checkpoint(checkpoint, f"{prefix}.mlp", output_dtype=output_dtype),
        )

    def forward(
        self,
        x: torch.Tensor,
        rope_rotation: torch.Tensor,
        shift_msa: torch.Tensor,
        scale_msa: torch.Tensor,
        gate_msa: torch.Tensor,
        shift_mlp: torch.Tensor,
        scale_mlp: torch.Tensor,
        gate_mlp: torch.Tensor,
        segments: list[tuple[int, int, int]],
    ) -> torch.Tensor:
        h = modulate_segments(rms_norm(x, self.norm1_weight, self.norm_eps), shift_msa, scale_msa, segments)
        x = gate_segments(x, self.attention(h, rope_rotation), gate_msa, segments)
        h = modulate_segments(rms_norm(x, self.norm2_weight, self.norm_eps), shift_mlp, scale_mlp, segments)
        return gate_segments(x, self.mlp(h), gate_mlp, segments)