"""Attention, positional features and native Inko transformer blocks."""

from typing import Tuple, Optional, Sequence
import math

import torch
import torch.nn as nn
import torch.nn.functional as F

from .varlen_attention import packed_attention

# Optional SDPA backend restriction for masked (padded-text) attention.
_SDPA_BACKENDS = None


class RMSNorm(nn.RMSNorm):
    """Fuse Ascend normalization while retaining FP32 arithmetic and gains."""

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if (
            x.device.type != "npu"
            or self.weight is None
            or len(self.normalized_shape) != 1
            or x.dtype not in (torch.float16, torch.bfloat16, torch.float32)
        ):
            return super().forward(x)
        import torch_npu

        # PyTorch's mixed-dtype RMSNorm promotes the input to FP32 and applies
        # the learned gain before casting back. Casting the gain to BF16 just
        # to use a fused kernel would change that computation.
        eps = self.eps if self.eps is not None else torch.finfo(x.dtype).eps
        y, _ = torch_npu.npu_rms_norm(x.float(), self.weight.float(), epsilon=eps)
        return y.to(x.dtype)


def sdpa_with_pad_mask(
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    attn_mask: Optional[torch.Tensor],
) -> torch.Tensor:
    """SDPA with an optional bool key-padding mask of shape [B, 1, 1, S_k].

    A bool mask forces PyTorch off the flash path; the math fallback then
    materializes B*H*S^2 attention weights per layer (saved for backward),
    which at seq ~2K is tens of GiB per layer and OOMs long-caption batches.
    We convert the mask to an additive bias in q.dtype and use fused kernels:
    cuDNN first in inference mode, memory-efficient attention otherwise.
    mask=None keeps the default (flash) path untouched.
    """
    if attn_mask is None:
        return F.scaled_dot_product_attention(q, k, v, dropout_p=0.0)

    return sdpa_with_bias(q, k, v, pad_bias_from_mask(attn_mask, q.dtype))


def set_sdpa_backends(names: Optional[Sequence[str]]) -> None:
    """Restrict SDPA backends for masked (padded-text) attention.

    ``None`` uses cuDNN first in inference mode and memory-efficient attention
    otherwise. An explicit list restricts attention to those backends, in order.
    Training keeps memory-efficient attention: on a
    single 4090 at 256p (micro-batch 16, 128-192 text tokens) cuDNN attention
    is ~30% slower than memory-efficient on the forward+backward training
    step, despite a faster forward-only kernel with an additive pad bias.
    """
    global _SDPA_BACKENDS
    if names is None:
        _SDPA_BACKENDS = None
        return
    from torch.nn.attention import SDPBackend

    resolved = []
    for name in names:
        member = getattr(SDPBackend, name.strip().upper(), None)
        if member is None:
            raise ValueError(f"unknown SDPA backend {name!r}")
        resolved.append(member)
    _SDPA_BACKENDS = tuple(resolved)


def _sdpa_backends():
    if _SDPA_BACKENDS is not None:
        return list(_SDPA_BACKENDS)
    from torch.nn.attention import SDPBackend

    if torch.is_inference_mode_enabled():
        return [
            SDPBackend.CUDNN_ATTENTION,
            SDPBackend.EFFICIENT_ATTENTION,
            SDPBackend.MATH,
        ]
    return [SDPBackend.EFFICIENT_ATTENTION, SDPBackend.MATH]


def sdpa_with_bias(
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    bias: torch.Tensor,
) -> torch.Tensor:
    """SDPA with a precomputed additive bias of shape [B, 1, 1, S_k]."""
    if q.device.type == "npu":
        # torch_npu's F.sdpa dispatch with a float additive bias materializes
        # per-head S^2 bias + P matrix per layer (measured: +2.8 GiB at
        # B8/S2304), which OOMs long-caption batches. npu_fusion_attention with
        # a bool (B,1,S,S) drop-mask (True = drop) is true flash attention:
        # +0.06 GiB saved for backward, maxdiff 2e-3 vs the float-bias path at
        # S=2304 bf16. Two gotchas, both measured: the scale argument defaults
        # to 1.0 (NOT 1/sqrt(D)) so it must be passed explicitly, and the mask
        # must be full (B,1,S,S) - aclnn rejects broadcastable (B,1,1,S).
        import torch_npu

        b, h, s_q, d = q.shape
        s_k = k.shape[2]
        # bias == -inf: identical to torch.isneginf for float tensors, but
        # aten.eq.Scalar has a torchair GE converter while aten.isneginf does
        # not (ascend-ta-bench3 NotImplementedError).
        drop = bias == float("-inf")
        mask = drop.expand(b, 1, s_q, s_k).contiguous()
        return torch_npu.npu_fusion_attention(
            q, k, v, h, "BNSD", atten_mask=mask, keep_prob=1.0, scale=d**-0.5
        )[0]
    from torch.nn.attention import sdpa_kernel

    backends = _sdpa_backends()
    if q.device.type != "cuda":
        from torch.nn.attention import SDPBackend

        backends = [SDPBackend.MATH]
    # sdpa_kernel only saves enabled backends. Preserve disabled backends'
    # positions too, so nested contexts do not reorder the caller's policy.
    priority = torch._C._get_sdp_priority_order()
    try:
        with sdpa_kernel(backends, set_priority=True):
            return F.scaled_dot_product_attention(q, k, v, attn_mask=bias, dropout_p=0.0)
    finally:
        torch._C._set_sdp_priority_order(priority)


def pad_bias_from_mask(attn_mask: torch.Tensor, dtype: torch.dtype) -> torch.Tensor:
    """Convert a bool keep-mask [B, 1, 1, S_k] into an additive SDPA bias."""
    bias = torch.zeros(attn_mask.shape, dtype=dtype, device=attn_mask.device)
    return bias.masked_fill(~attn_mask, float("-inf"))


class TimestepEmbeddings(nn.Module):
    """Sinusoidal features; scaling applies here, never to the flow path."""

    def __init__(self, hidden_size: int):
        super().__init__()
        self.hidden_size = hidden_size
        self.register_buffer("factor", torch.tensor(1000, dtype=torch.int64))

    def _load_from_state_dict(
        self,
        state_dict,
        prefix,
        local_metadata,
        strict,
        missing_keys,
        unexpected_keys,
        error_msgs,
    ):
        factor = state_dict.get(prefix + "factor")
        if factor is None or factor.item() != 1000:
            error_msgs.append(f"{prefix}expected timestep feature factor 1000")
        super()._load_from_state_dict(
            state_dict,
            prefix,
            local_metadata,
            strict,
            missing_keys,
            unexpected_keys,
            error_msgs,
        )

    def forward(self, timesteps: torch.Tensor) -> torch.Tensor:
        """
        Args:
            timesteps: (B,)
        Returns:
            embedding: (B, N)
        """
        half = self.hidden_size // 2
        exponent = (
            -math.log(10000)
            * torch.arange(
                start=0, end=half, dtype=torch.float32, device=timesteps.device
            )
            / half
        )

        emb = torch.exp(exponent)
        times = timesteps.float() * 1000
        emb = times.unsqueeze(1) * emb.unsqueeze(0)

        return torch.cat([torch.sin(emb), torch.cos(emb)], dim=-1)


def apply_rotary_emb(x: torch.Tensor, freqs_cis: torch.Tensor) -> torch.Tensor:
    """
    Apply rotary embeddings
    Args:
        x: Input tensor [B, S, H, D]
        freqs_cis: Complex frequencies [S, D/2] or [B, S, D/2]
    Returns:
        Tensor with rotary embeddings applied [B, S, H, D]
    """
    # Reshape x to [B, S, H, D/2, 2] and view as complex
    x_complex = torch.view_as_complex(x.float().reshape(*x.shape[:-1], -1, 2))

    # Reshape for broadcasting: [S, D/2] -> [1, S, 1,  D/2]
    if freqs_cis.ndim == 2:
        freqs_cis = freqs_cis.unsqueeze(0)
    freqs_cis = freqs_cis.unsqueeze(2)

    # Apply rotation
    x_rotated = x_complex * freqs_cis

    # View as real and flatten the last two dimensions
    x_out = torch.view_as_real(x_rotated).flatten(3)

    return x_out.type_as(x)


def apply_rotary_emb_real(x: torch.Tensor, freqs_cis: torch.Tensor) -> torch.Tensor:
    """Equivalent FP32 rotation expressed without complex multiplication.

    This lets Inductor fuse surrounding casts/layout operations. Compiled
    rounding may differ; the complex implementation remains the default.
    """
    pairs = x.float().reshape(*x.shape[:-1], -1, 2)
    real, imag = pairs.unbind(-1)
    freqs_real = torch.view_as_real(freqs_cis)
    if freqs_real.ndim == 3:
        freqs_real = freqs_real.unsqueeze(0)
    cosine, sine = freqs_real.unsqueeze(2).unbind(-1)
    rotated = torch.stack(
        (real * cosine - imag * sine, real * sine + imag * cosine), dim=-1
    )
    return rotated.flatten(3).type_as(x)


def apply_rotary_emb_realfreq(
    x: torch.Tensor, freqs_real: torch.Tensor
) -> torch.Tensor:
    """apply_rotary_emb_real taking an already-real frequency tensor.

    freqs_real is [S, D/2, 2] or [B, S, D/2, 2] float32 (cos, sin pairs), the view_as_real of
    the complex table, materialized OUTSIDE compiled regions so no
    complex-typed tensor crosses the graph boundary (torchair/GE has no
    complex dtype support).
    """
    pairs = x.float().reshape(*x.shape[:-1], -1, 2)
    real, imag = pairs.unbind(-1)
    if freqs_real.ndim == 3:
        freqs_real = freqs_real.unsqueeze(0)
    cosine, sine = freqs_real.unsqueeze(2).unbind(-1)
    rotated = torch.stack(
        (real * cosine - imag * sine, real * sine + imag * cosine), dim=-1
    )
    return rotated.flatten(3).type_as(x)


def set_real_rope(model: nn.Module, enabled: bool) -> None:
    """Set an instance-local execution policy before compiling the model.

    This is not checkpoint state: loading weights must not override the launch
    policy. Other live models and the complex frequency tables are untouched.
    """
    if type(enabled) is not bool:
        raise ValueError("real RoPE policy must be boolean")
    for module in model.modules():
        if isinstance(module, (DoubleStreamAttention, SingleStreamAttention)):
            module.real_rope = enabled


class MSRoPE(nn.Module):
    """Image/text RoPE with automatic Dy-YaRN inference scaling.

    The 896p training buckets reach 1184 pixels on either axis, or 74 patches
    at 16 pixels per patch. Training always uses the original frequencies.
    DyPE follows https://arxiv.org/html/2510.20766v2, equations 9 and 13–17.
    Both axes share one scaling factor to preserve the canvas aspect ratio
    in phase space.
    """

    def __init__(
        self,
        theta: int = 10000,
        axes_dim: list = [64, 64],
        reference_hw: tuple[int, int] = (74, 74),
        centered: bool = False,
        decay_exponent: float = 2.0,
        yarn_alpha: float = 1.0,
        yarn_beta: float = 32.0,
    ):
        super().__init__()
        self.theta = theta
        self.axes_dim = axes_dim
        self.reference_hw = reference_hw
        self.decay_exponent = decay_exponent
        self.yarn_alpha = yarn_alpha
        self.yarn_beta = yarn_beta
        # centered: image coords symmetric around 0 ([-(n-1)/2, (n-1)/2]) instead of
        # 0-based from the corner. Text starts at max(height, width) on both axes.
        self.centered = centered

        # Precompute frequency tables for both spatial axes
        pos_index = torch.arange(2560)  # Initial table capacity; expands as needed.
        # Store as real to avoid safetensors issues with complex numbers
        self.register_buffer(
            "pos_freqs", torch.view_as_real(self._build_frequency_table(pos_index))
        )

        # Retain only the latest geometry/mode on each instance.
        self._cache = {}

    def needs_time(self, img_hw: tuple[int, int]) -> bool:
        return not self.training and any(
            size > reference for size, reference in zip(img_hw, self.reference_hw)
        )

    def _apply(self, fn, recurse=True):
        super()._apply(fn, recurse=recurse)
        if self.pos_freqs.dtype != torch.float32:
            # Rebuild instead of promoting already-rounded BF16/FP16 rotations.
            self.pos_freqs = torch.view_as_real(
                self._build_frequency_table(torch.arange(self.pos_freqs.shape[0]))
            ).to(self.pos_freqs.device)
        self._cache.clear()
        return self

    def rope_params(
        self,
        index: torch.Tensor,
        dim: int,
        scaling_factor: float = 1.0,
        *,
        reference: int,
        kappa: float = 1.0,
    ) -> torch.Tensor:
        """
        Generate RoPE parameters for given dimensions
        Args:
            index: Position indices [S]
            dim: Embedding dimension for this axis
            scaling_factor: Spatial extension along this axis
        Returns:
            Complex frequency tensor [S, dim/2]
        """
        assert dim % 2 == 0

        # Compute frequency components: 1/ theta^(2i/dim)
        powers = torch.arange(0, dim, 2, dtype=torch.float32, device=index.device).div(
            dim
        )
        inverse_freqs = 1.0 / torch.pow(self.theta, powers)
        if scaling_factor > 1.0 and kappa > 0.0:
            rotations = reference * inverse_freqs / (2 * math.pi)
            ramp = (
                (rotations - self.yarn_alpha * kappa)
                / ((self.yarn_beta - self.yarn_alpha) * kappa)
            ).clamp(0, 1)
            inverse_freqs = inverse_freqs * (ramp + (1 - ramp) / scaling_factor)
        freqs = torch.outer(index.float(), inverse_freqs)

        # Convert to complex polar form: e^(i*freqs)
        return torch.polar(torch.ones_like(freqs), freqs)

    def _build_frequency_table(
        self,
        pos_index: torch.Tensor,
        scaling_factors: tuple[float, float] = (1.0, 1.0),
        *,
        kappa: float = 1.0,
    ) -> torch.Tensor:
        """Build frequency table for both height and width axes"""
        height_freqs = self.rope_params(
            pos_index,
            self.axes_dim[0],
            scaling_factors[0],
            kappa=kappa,
            reference=self.reference_hw[0],
        )
        width_freqs = self.rope_params(
            pos_index,
            self.axes_dim[1],
            scaling_factors[1],
            kappa=kappa,
            reference=self.reference_hw[1],
        )

        # Concatenate height and width frequencies: [S, (height_dim+width_dim)/2]
        return torch.cat([height_freqs, width_freqs], dim=1)

    def forward(
        self,
        img_hw: tuple[int, int],
        txt_seq_len: int,
        device,
        *,
        time: tuple[float, ...] | None = None,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """
        Compute frequency tensors for image and text sequences
        Args:
            img_hw: Image patch-grid dimensions (height, width)
            txt_seq_len: Length of text sequence
            device: Device for tensor placement
            time: Flow time per image (0 noise, 1 image); required for active DyPE
        Returns:
            Image/text frequencies, with a batch dimension when times differ
        """
        height, width = img_hw
        device = torch.device(device)
        scaling_factors = (1.0, 1.0)
        if not self.training:
            shared = max(
                max(1.0, size / reference)
                for size, reference in zip(img_hw, self.reference_hw)
            )
            scaling_factors = (shared, shared)
        scaled = scaling_factors != (1.0, 1.0)
        kappas = (1.0,)
        if self.needs_time(img_hw):
            if time is None:
                raise ValueError("dynamic RoPE requires the current flow time")
            # Inko runs from t=0 noise to t=1 image, opposite to the paper.
            kappas = tuple((1 - t) ** self.decay_exponent for t in time)
            if len(set(kappas)) == 1:
                kappas = kappas[:1]

        # Keep the checkpoint's base table unscaled.
        max_needed = max(height, width) + txt_seq_len
        current_max = self.pos_freqs.shape[0]
        if max_needed > current_max:
            new_max = max(max_needed, current_max * 2)
            self.pos_freqs = torch.view_as_real(
                self._build_frequency_table(torch.arange(new_max))
            )
        if self.pos_freqs.device != device:
            self.pos_freqs = self.pos_freqs.to(device)

        cache_key = (
            height,
            width,
            txt_seq_len,
            scaling_factors,
            kappas,
            self.reference_hw,
            self.yarn_alpha,
            self.yarn_beta,
            self.centered,
            device,
        )
        if cache_key in self._cache:
            return self._cache[cache_key]

        images, texts = [], []
        max_img_pos = max(height, width)
        for kappa in kappas:
            images.append(
                self._compute_image_freqs(
                    height, width, scaling_factors, kappa=kappa
                ).to(device)
            )
            # Share the image basis without relocating the text diagonal.
            texts.append(
                torch.view_as_complex(self.pos_freqs)[max_img_pos:max_needed]
                if not scaled
                else self._build_frequency_table(
                    torch.arange(max_img_pos, max_needed),
                    scaling_factors,
                    kappa=kappa,
                ).to(device)
            )
        img_freqs = images[0] if len(images) == 1 else torch.stack(images)
        txt_freqs = texts[0] if len(texts) == 1 else torch.stack(texts)
        if scaled:
            # Eq. 9 multiplies logits by tau. Multiplying BOTH Q/K rotations by
            # sqrt(tau) applies it in every attention backend and RoPE layout.
            amplitude = math.sqrt(1 + 0.1 * math.log(max(scaling_factors)))
            img_freqs = img_freqs * amplitude
            txt_freqs = txt_freqs * amplitude

        self._cache.clear()
        self._cache[cache_key] = (img_freqs, txt_freqs)
        return img_freqs, txt_freqs

    def prepare_freqs(
        self,
        img_hw: tuple[int, int],
        txt_seq_len: int,
        device,
        *,
        time: tuple[float, ...] | None = None,
    ) -> torch.Tensor:
        """Concatenated [S_img + S_txt, D/2] freqs, optionally batched.

        Identical to applying ``forward()`` and concatenating the two halves;
        hoisting it out of the block loop removes one table lookup and two
        concatenations per layer.
        """
        img_freqs, txt_freqs = self.forward(img_hw, txt_seq_len, device, time=time)
        return torch.cat([img_freqs, txt_freqs], dim=-2)

    def _compute_image_freqs(
        self,
        height: int,
        width: int,
        scaling_factors: tuple[float, float],
        *,
        kappa: float = 1.0,
    ) -> torch.Tensor:
        """
        Compute frequency tensor for 2D image grid
        Args:
            height: Height of image
            width: Width of image
        Returns:
            Frequency tensor [height*width, total_dim] (complex)
        """
        h_dim, w_dim = self.axes_dim
        h_idx = torch.arange(height, dtype=torch.float32)
        w_idx = torch.arange(width, dtype=torch.float32)
        if self.centered:
            # Symmetric (possibly half-integer) coordinates around the origin.
            h_idx = h_idx - (height - 1) / 2.0
            w_idx = w_idx - (width - 1) / 2.0
        h_freqs = torch.view_as_real(
            self.rope_params(
                h_idx,
                h_dim,
                scaling_factors[0],
                kappa=kappa,
                reference=self.reference_hw[0],
            )
        )
        w_freqs = torch.view_as_real(
            self.rope_params(
                w_idx,
                w_dim,
                scaling_factors[1],
                kappa=kappa,
                reference=self.reference_hw[1],
            )
        )

        # Broadcast to create the grid
        # Expand to [H, W, dim, 2]
        h_freqs_grid = h_freqs.unsqueeze(1).expand(
            height, width, -1, -1
        )  # [height, width, h_dim//2, 2]
        w_freqs_grid = w_freqs.unsqueeze(0).expand(
            height, width, -1, -1
        )  # [height, width, w_dim//2, 2]

        # Concatenate
        freqs = torch.cat(
            [h_freqs_grid, w_freqs_grid], dim=2
        )  # [height, width, (h_dim+w_dim)//2, 2]

        # Flatten spatial dimensions
        freqs = freqs.flatten(0, 1)  # [height*width, (h_dim+w_dim)//2, 2]

        # Convert back to complex
        freqs = torch.view_as_complex(freqs)

        return freqs


class GatedFeedForward(nn.Module):
    def __init__(self, dim: int, hidden_dim: int, dropout: float = 0.0):
        super().__init__()
        self.up_proj = nn.Linear(dim, hidden_dim * 2)
        self.down_proj = nn.Linear(hidden_dim, dim)
        self.dropout = nn.Dropout(dropout)
        self.dropout_out = nn.Dropout(dropout)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.up_proj(x)
        if x.device.type == "npu":
            import torch_npu

            x = torch_npu.npu_swiglu(x, dim=-1)
        else:
            x_gated, x_linear = x.chunk(2, dim=-1)
            x = F.silu(x_gated) * x_linear
        x = self.dropout(x)
        x = self.down_proj(x)
        return self.dropout_out(x)


def modulate(x: torch.Tensor, shift: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
    return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)


class DoubleStreamAttention(nn.Module):
    def __init__(
        self,
        dim: int,
        num_heads: int,
        qkv_bias: bool = True,
        rope_theta: int = 10000,
        rope_axes_dim: list = [64, 64],
        rope_centered: bool = False,
    ):
        super().__init__()
        self.num_heads = num_heads
        self.head_dim = dim // num_heads
        self.scale = self.head_dim**-0.5
        self.real_rope = False

        self.qkv_img = nn.Linear(dim, dim * 3, bias=qkv_bias)
        self.qkv_txt = nn.Linear(dim, dim * 3, bias=qkv_bias)

        self.q_norm_img = RMSNorm(self.head_dim, eps=1e-6)
        self.k_norm_img = RMSNorm(self.head_dim, eps=1e-6)
        self.q_norm_txt = RMSNorm(self.head_dim, eps=1e-6)
        self.k_norm_txt = RMSNorm(self.head_dim, eps=1e-6)

        self.rope = MSRoPE(
            theta=rope_theta,
            axes_dim=rope_axes_dim,
            centered=rope_centered,
        )

        self.proj_img = nn.Linear(dim, dim)
        self.proj_txt = nn.Linear(dim, dim)

    def forward(
        self,
        img_tokens: torch.Tensor,
        txt_tokens: torch.Tensor,
        img_hw: Tuple[int, int],
        txt_seq_len: int,
        txt_attention_mask: Optional[torch.Tensor] = None,
        attn_metadata=None,
        rope_freqs: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
        rope_time: tuple[float, ...] | None = None,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        B, S_img, C = img_tokens.shape
        _, S_txt, _ = txt_tokens.shape

        # Image QKV
        qkv_img = (
            self.qkv_img(img_tokens)
            .reshape(B, S_img, 3, self.num_heads, self.head_dim)
            .permute(2, 0, 3, 1, 4)
        )
        q_img, k_img, v_img = qkv_img.unbind(0)  # [B, H, S, D]

        # Text QKV
        qkv_txt = (
            self.qkv_txt(txt_tokens)
            .reshape(B, S_txt, 3, self.num_heads, self.head_dim)
            .permute(2, 0, 3, 1, 4)
        )
        q_txt, k_txt, v_txt = qkv_txt.unbind(0)  # [B, H, S, D]

        # QK Norm
        q_img = self.q_norm_img(q_img)
        k_img = self.k_norm_img(k_img)
        q_txt = self.q_norm_txt(q_txt)
        k_txt = self.k_norm_txt(k_txt)

        # RoPE
        img_freqs, txt_freqs = (
            self.rope(img_hw, txt_seq_len, img_tokens.device, time=rope_time)
            if rope_freqs is None
            else rope_freqs
        )

        # Apply RoPE (need to transpose to [B, S, H, D] for apply_rotary_emb)
        rotate = apply_rotary_emb_real if self.real_rope else apply_rotary_emb
        q_img = rotate(q_img.transpose(1, 2), img_freqs).transpose(1, 2)
        k_img = rotate(k_img.transpose(1, 2), img_freqs).transpose(1, 2)

        q_txt = rotate(q_txt.transpose(1, 2), txt_freqs).transpose(1, 2)
        k_txt = rotate(k_txt.transpose(1, 2), txt_freqs).transpose(1, 2)

        # Concat
        q = torch.cat([q_img, q_txt], dim=2)
        k = torch.cat([k_img, k_txt], dim=2)
        v = torch.cat([v_img, v_txt], dim=2)

        # Prepare attention mask
        attn_mask = None
        if txt_attention_mask is not None and attn_metadata is None:
            # Image tokens always have attention
            img_mask = torch.ones(
                B,
                S_img,
                device=txt_attention_mask.device,
                dtype=txt_attention_mask.dtype,
            )
            full_mask = torch.cat([img_mask, txt_attention_mask], dim=1)

            # Convert to attention mask format [B, 1, 1, S_total]
            # Mask: 1 (valid) -> True (attend), 0 (padding) -> False (ignore)
            attn_mask = (full_mask > 0).unsqueeze(1).unsqueeze(2)
            attn_mask = attn_mask.to(dtype=torch.bool)

        # Attention with mask
        x = (
            packed_attention(q, k, v, attn_metadata)
            if attn_metadata is not None
            else sdpa_with_pad_mask(q, k, v, attn_mask)
        )

        # Split
        x_img = x[:, :, :S_img, :]
        x_txt = x[:, :, S_img:, :]

        # Reshape and Project
        x_img = x_img.transpose(1, 2).reshape(B, S_img, C)
        x_txt = x_txt.transpose(1, 2).reshape(B, S_txt, C)

        x_img = self.proj_img(x_img)
        x_txt = self.proj_txt(x_txt)

        return x_img, x_txt


class DoubleStreamDiTBlock(nn.Module):
    def __init__(
        self,
        dim: int,
        num_heads: int,
        c_dim: int,
        mlp_ratio: float = 4.0,
        rope_theta: int = 10000,
        rope_axes_dim: list = [64, 64],
    ):
        super().__init__()
        self.dim = dim
        self.num_heads = num_heads
        self.mlp_ratio = mlp_ratio

        # Normalize branch outputs before learned residual gates.
        self.norm_msa_img_out = RMSNorm(dim, eps=1e-6)
        self.norm_msa_txt_out = RMSNorm(dim, eps=1e-6)
        self.norm_mlp_img_out = RMSNorm(dim, eps=1e-6)
        self.norm_mlp_txt_out = RMSNorm(dim, eps=1e-6)

        # Modulation
        self.modulation_img = nn.Sequential(
            nn.SiLU(), nn.Linear(c_dim, 6 * dim, bias=True)
        )
        self.modulation_txt = nn.Sequential(
            nn.SiLU(), nn.Linear(c_dim, 6 * dim, bias=True)
        )

        # Attention
        self.norm1_img = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
        self.norm1_txt = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
        self.attn = DoubleStreamAttention(
            dim,
            num_heads,
            True,
            rope_theta,
            rope_axes_dim,
            rope_centered=True,
        )

        # MLP
        self.norm2_img = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
        self.norm2_txt = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
        mlp_hidden_dim = int(dim * mlp_ratio)

        self.mlp_img = GatedFeedForward(dim, mlp_hidden_dim)
        self.mlp_txt = GatedFeedForward(dim, mlp_hidden_dim)

    def forward(
        self,
        img_tokens: torch.Tensor,
        txt_tokens: torch.Tensor,
        c: torch.Tensor,
        img_hw: Tuple[int, int],
        txt_seq_len: int,
        txt_attention_mask: Optional[torch.Tensor] = None,
        attn_metadata=None,
        rope_freqs: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
        rope_time: tuple[float, ...] | None = None,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        # Modulation
        # c: [B, c_dim]
        (
            shift_msa_img,
            scale_msa_img,
            gate_msa_img,
            shift_mlp_img,
            scale_mlp_img,
            gate_mlp_img,
        ) = self.modulation_img(c).chunk(6, dim=1)
        (
            shift_msa_txt,
            scale_msa_txt,
            gate_msa_txt,
            shift_mlp_txt,
            scale_mlp_txt,
            gate_mlp_txt,
        ) = self.modulation_txt(c).chunk(6, dim=1)

        # 1. Attention Block
        img_norm = modulate(self.norm1_img(img_tokens), shift_msa_img, scale_msa_img)
        txt_norm = modulate(self.norm1_txt(txt_tokens), shift_msa_txt, scale_msa_txt)

        img_attn, txt_attn = self.attn(
            img_norm,
            txt_norm,
            img_hw,
            txt_seq_len,
            txt_attention_mask,
            attn_metadata=attn_metadata,
            rope_freqs=rope_freqs,
            rope_time=rope_time,
        )

        img_attn = self.norm_msa_img_out(img_attn)
        txt_attn = self.norm_msa_txt_out(txt_attn)
        img_tokens = img_tokens + gate_msa_img.unsqueeze(1) * img_attn
        txt_tokens = txt_tokens + gate_msa_txt.unsqueeze(1) * txt_attn

        # 2. MLP Block
        img_norm = modulate(self.norm2_img(img_tokens), shift_mlp_img, scale_mlp_img)
        txt_norm = modulate(self.norm2_txt(txt_tokens), shift_mlp_txt, scale_mlp_txt)

        mlp_img_out = self.mlp_img(img_norm)
        mlp_txt_out = self.mlp_txt(txt_norm)
        mlp_img_out = self.norm_mlp_img_out(mlp_img_out)
        mlp_txt_out = self.norm_mlp_txt_out(mlp_txt_out)
        img_tokens = img_tokens + gate_mlp_img.unsqueeze(1) * mlp_img_out
        txt_tokens = txt_tokens + gate_mlp_txt.unsqueeze(1) * mlp_txt_out

        return img_tokens, txt_tokens

    def initialize_weights(self):
        # AdaLN-zero: Initialize modulation MLP to zero
        nn.init.constant_(self.modulation_img[-1].weight, 0)
        nn.init.constant_(self.modulation_img[-1].bias, 0)
        nn.init.constant_(self.modulation_txt[-1].weight, 0)
        nn.init.constant_(self.modulation_txt[-1].bias, 0)


class SingleStreamAttention(nn.Module):
    def __init__(
        self,
        dim: int,
        num_heads: int,
        qkv_bias: bool = True,
        rope_theta: int = 10000,
        rope_axes_dim: list = [64, 64],
        rope_centered: bool = False,
    ):
        super().__init__()
        self.num_heads = num_heads
        self.head_dim = dim // num_heads
        self.scale = self.head_dim**-0.5
        self.real_rope = False

        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
        self.q_norm = RMSNorm(self.head_dim, eps=1e-6)
        self.k_norm = RMSNorm(self.head_dim, eps=1e-6)

        self.rope = MSRoPE(
            theta=rope_theta,
            axes_dim=rope_axes_dim,
            centered=rope_centered,
        )

        self.proj = nn.Linear(dim, dim)

    def forward(
        self,
        x: torch.Tensor,
        img_hw: Tuple[int, int],
        txt_seq_len: int,
        txt_attention_mask: Optional[torch.Tensor] = None,
        rope_freqs: Optional[torch.Tensor] = None,
        attn_bias: Optional[torch.Tensor] = None,
        attn_metadata=None,
        rope_time: tuple[float, ...] | None = None,
    ) -> torch.Tensor:
        B, S, C = x.shape

        qkv = (
            self.qkv(x)
            .reshape(B, S, 3, self.num_heads, self.head_dim)
            .permute(2, 0, 3, 1, 4)
        )
        q, k, v = qkv.unbind(0)

        q = self.q_norm(q)
        k = self.k_norm(k)

        if rope_freqs is not None:
            # Frequencies already cover [img, txt] in that order, so RoPE is a
            # single elementwise pass over the concatenated sequence. Under
            # the real-rope policy the hoisted table arrives in real
            # [S, D/2, 2] layout (converted once at model level, outside any
            # compiled block); otherwise it is complex.
            rotate = apply_rotary_emb_realfreq if self.real_rope else apply_rotary_emb
            q = rotate(q.transpose(1, 2), rope_freqs).transpose(1, 2)
            k = rotate(k.transpose(1, 2), rope_freqs).transpose(1, 2)
        else:
            rotate = apply_rotary_emb_real if self.real_rope else apply_rotary_emb
            # Need to apply different RoPE to img and txt parts
            # Assuming x is [img, txt]
            S_img = img_hw[0] * img_hw[1]

            q_img = q[:, :, :S_img, :]
            k_img = k[:, :, :S_img, :]
            q_txt = q[:, :, S_img:, :]
            k_txt = k[:, :, S_img:, :]

            img_freqs, txt_freqs = self.rope(
                img_hw, txt_seq_len, x.device, time=rope_time
            )

            q_img = rotate(q_img.transpose(1, 2), img_freqs).transpose(1, 2)
            k_img = rotate(k_img.transpose(1, 2), img_freqs).transpose(1, 2)
            q_txt = rotate(q_txt.transpose(1, 2), txt_freqs).transpose(1, 2)
            k_txt = rotate(k_txt.transpose(1, 2), txt_freqs).transpose(1, 2)

            q = torch.cat([q_img, q_txt], dim=2)
            k = torch.cat([k_img, k_txt], dim=2)

        if attn_metadata is not None:
            x = packed_attention(q, k, v, attn_metadata)
            return self.proj(x.transpose(1, 2).reshape(B, S, C))

        if attn_bias is not None:
            x = sdpa_with_bias(q, k, v, attn_bias)
            x = x.transpose(1, 2).reshape(B, S, C)
            return self.proj(x)

        # Prepare attention mask
        attn_mask = None
        if txt_attention_mask is not None:
            # Image tokens always have attention
            img_mask = torch.ones(
                B,
                S_img,
                device=txt_attention_mask.device,
                dtype=txt_attention_mask.dtype,
            )
            full_mask = torch.cat([img_mask, txt_attention_mask], dim=1)

            # Convert to attention mask format [B, 1, 1, S]
            # Mask : 1 (valid) -> True (attend), 0 (padding) -> False (ignore)
            attn_mask = (full_mask > 0).unsqueeze(1).unsqueeze(2)
            attn_mask = attn_mask.to(dtype=torch.bool)

        # Attention with mask
        x = sdpa_with_pad_mask(q, k, v, attn_mask)

        x = x.transpose(1, 2).reshape(B, S, C)

        x = self.proj(x)

        return x


class SingleStreamDiTBlock(nn.Module):
    def __init__(
        self,
        dim: int,
        num_heads: int,
        c_dim: int,
        mlp_ratio: float = 4.0,
        rope_theta: int = 10000,
        rope_axes_dim: list = [64, 64],
    ):
        super().__init__()
        self.dim = dim
        self.num_heads = num_heads
        self.mlp_ratio = mlp_ratio

        # Modulation
        self.modulation = nn.Sequential(nn.SiLU(), nn.Linear(c_dim, 3 * dim, bias=True))

        # Normalize branch outputs before the residual add. Learned gates and
        # affine norm gains remain outside this control of branch amplitude.
        self.norm_msa_out = RMSNorm(dim, eps=1e-6)
        self.norm_mlp_out = RMSNorm(dim, eps=1e-6)

        self.norm1 = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
        self.attn = SingleStreamAttention(
            dim,
            num_heads,
            True,
            rope_theta,
            rope_axes_dim,
            rope_centered=True,
        )

        # MLP
        self.norm2 = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
        mlp_hidden_dim = int(dim * mlp_ratio)

        self.mlp = GatedFeedForward(dim, mlp_hidden_dim)

    def forward(
        self,
        img_tokens: torch.Tensor,
        txt_tokens: torch.Tensor,
        c: torch.Tensor,
        img_hw: Tuple[int, int],
        txt_seq_len: int,
        txt_attention_mask: Optional[torch.Tensor] = None,
        rope_freqs: Optional[torch.Tensor] = None,
        attn_bias: Optional[torch.Tensor] = None,
        attn_metadata=None,
        rope_time: tuple[float, ...] | None = None,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        # Modulation
        # c: [B, c_dim]
        params = self.modulation(c)
        shift, scale, gate = params.chunk(3, dim=1)
        shift_msa = shift_mlp = shift
        scale_msa = scale_mlp = scale
        gate_msa = gate_mlp = gate

        # Concatenate for shared MSA/MLP processing
        S_img = img_tokens.shape[1]
        S_txt = txt_tokens.shape[1]
        x = torch.cat([img_tokens, txt_tokens], dim=1)

        # 1. Attention Block
        x_norm = modulate(self.norm1(x), shift_msa, scale_msa)

        x_attn = self.attn(
            x_norm,
            img_hw,
            txt_seq_len,
            txt_attention_mask,
            rope_freqs=rope_freqs,
            attn_bias=attn_bias,
            attn_metadata=attn_metadata,
            rope_time=rope_time,
        )

        x_attn = self.norm_msa_out(x_attn)
        x = x + gate_msa.unsqueeze(1) * x_attn

        # 2. MLP Block
        x_norm = modulate(self.norm2(x), shift_mlp, scale_mlp)
        mlp_out = self.mlp(x_norm)
        mlp_out = self.norm_mlp_out(mlp_out)
        x = x + gate_mlp.unsqueeze(1) * mlp_out

        img_tokens, txt_tokens = x.split([S_img, S_txt], dim=1)

        return img_tokens, txt_tokens

    def initialize_weights(self):
        # AdaLN-zero: Initialize modulation MLP to zero
        nn.init.constant_(self.modulation[-1].weight, 0)
        nn.init.constant_(self.modulation[-1].bias, 0)