"""Direct 50-block H3 denoiser backbone with standalone Sage3 attention."""

import torch
from torch import nn

from .adaln import H3CurveAdaLN
from .block import H3DiTBlock
from .checkpoint import H3Checkpoint
from .rope import h3_rope_rotation


class H3DenoiserBackbone(nn.Module):
    """Execute H3 transformer blocks over an already packed Ref2VA hidden sequence."""

    def __init__(self, blocks: list[H3DiTBlock], adaln: list[H3CurveAdaLN], inv_freq: torch.Tensor):
        super().__init__()
        if len(blocks) != 50 or len(adaln) != 50:
            raise ValueError("The released H3 denoiser has exactly 50 transformer blocks.")
        self.blocks = nn.ModuleList(blocks)
        self.adaln = nn.ModuleList(adaln)
        self.register_buffer("inv_freq", inv_freq.to(torch.float32), persistent=False)

    @classmethod
    def from_checkpoint(cls, checkpoint: H3Checkpoint, *, output_dtype=torch.bfloat16, attention_backend: str = "sage2"):
        return cls(
            [H3DiTBlock.from_checkpoint(checkpoint, index, output_dtype=output_dtype, attention_backend=attention_backend) for index in range(50)],
            [H3CurveAdaLN.from_checkpoint(checkpoint, f"blocks.{index}.adaln_proj") for index in range(50)],
            checkpoint.tensor("rope.inv_freq", dtype=torch.float32),
        )

    def forward(
        self,
        hidden: torch.Tensor,
        timesteps: torch.Tensor,
        position_ids: torch.Tensor,
        segments: list[tuple[int, int, int]],
    ) -> torch.Tensor:
        rotation = h3_rope_rotation(position_ids.to(hidden.device), self.inv_freq, hidden.dtype)
        for block, adaln in zip(self.blocks, self.adaln, strict=True):
            hidden = block(hidden, rotation, *adaln(timesteps), segments)
        return hidden