"""Curve-form AdaLN used by the pruned H3 NVFP4 checkpoint."""

import torch
import torch.nn.functional as functional
from torch import nn

from .checkpoint import H3Checkpoint


class H3CurveAdaLN(nn.Module):
    """Interpolate the H3 timestep curve and emit six modality-specific AdaLN tensors."""

    def __init__(self, curve_table: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor, *, hidden_size: int = 5376):
        super().__init__()
        if curve_table.shape != (1025, 8):
            raise ValueError(f"Unexpected H3 AdaLN curve shape: {tuple(curve_table.shape)}")
        if weight.shape != (18 * hidden_size, curve_table.shape[1]) or bias.shape != (18 * hidden_size,):
            raise ValueError("Unexpected H3 AdaLN projection dimensions.")
        self.hidden_size = hidden_size
        self.register_buffer("curve_table", curve_table.to(torch.float32), persistent=False)
        self.register_buffer("weight", weight.to(torch.float32), persistent=False)
        self.register_buffer("bias", bias.to(torch.float32), persistent=False)

    @classmethod
    def from_checkpoint(cls, checkpoint: H3Checkpoint, prefix: str):
        return cls(
            checkpoint.tensor("adaln_t_table", dtype=torch.float32),
            checkpoint.tensor(f"{prefix}.linear.weight", dtype=torch.bfloat16),
            checkpoint.tensor(f"{prefix}.linear.bias", dtype=torch.bfloat16),
        )

    def forward(self, timesteps: torch.Tensor) -> tuple[torch.Tensor, ...]:
        """Return shift/scale/gate tensors ordered as MSA then MLP for all modalities."""
        table = self.curve_table
        position = timesteps.float().clamp(0, 1) * (table.shape[0] - 1)
        lower = position.floor().long().clamp(max=table.shape[0] - 2)
        embedding = torch.lerp(table[lower], table[lower + 1], (position - lower).unsqueeze(1))
        values = functional.linear(embedding, self.weight, self.bias)
        return values.view(values.shape[0] * 3, 6 * self.hidden_size).chunk(6, dim=-1)