"""Blackwell-native MiniMax-H3 transformer for the pruned ComfyUI NVFP4 checkpoint.

The public diffusers checkpoint spends 13.04B of its 33.12B parameters on per-block
AdaLN projections.  ComfyUI's pruned checkpoint replaces those projections with an
interpolated 1025-point timestep curve, fuses Q/K/V, and stores the four large linear
layers in every block as NVFP4.  This adapter keeps diffusers' packed-sequence contract
so the rest of the split Space (schedulers, VAEs and remote conditioner) stays unchanged.

The kernel/layout conventions follow ComfyUI's Apache-2.0 implementation:
https://github.com/Comfy-Org/ComfyUI/blob/master/comfy/ldm/minimax/model.py
"""

from __future__ import annotations

import json
import math
import os
from types import SimpleNamespace

import torch
import torch.nn as nn
import torch.nn.functional as F
import comfy_kitchen as kitchen
from comfy_kitchen.tensor import QuantizedTensor, TensorCoreNVFP4Layout
from diffusers.models.attention_dispatch import dispatch_attention_fn

try:
    import triton
    import triton.language as tl
except ImportError:  # PyTorch CUDA wheels include Triton; retain a portable fallback for source inspection/tests.
    triton = None
    tl = None


NVFP4_REPO = os.environ.get("H3_NVFP4_REPO", "lilcheaty/MiniMax-H3-NVFP4")
NVFP4_FILE = os.environ.get("H3_NVFP4_FILE", "minimax_h3_fl2va_pruned_nvfp4.safetensors")

HIDDEN = 5376
HEADS = 56
HEAD_DIM = 128
FFN = 14336
TEXT_DIM = 5120
TIME_DIM = 8
VIDEO_DIM = 24 * 1 * 2 * 2
AUDIO_DIM = 32
LAYERS = 50
REFINER_LAYERS = 2
EPS = 1e-5

# EasyCache is the conservative profile.  The Ultra Fast profile uses a bounded linear residual forecast:
# three exact warmup evaluations, at most three forecasts in a row, and two exact tail evaluations.  Unlike blind
# output reuse, forecasting follows the local denoising trajectory while making the amount of saved work predictable.
EASYCACHE_THRESHOLD = max(0.0, float(os.environ.get("H3_EASYCACHE_THRESHOLD", "0.10")))
EASYCACHE_START = min(1.0, max(0.0, float(os.environ.get("H3_EASYCACHE_START", "0.15"))))
EASYCACHE_END = min(1.0, max(EASYCACHE_START, float(os.environ.get("H3_EASYCACHE_END", "0.95"))))
EASYCACHE_SUBSAMPLE = max(1, int(os.environ.get("H3_EASYCACHE_SUBSAMPLE", "8")))
FIRST_BLOCK_THRESHOLD = max(0.0, float(os.environ.get("H3_FIRST_BLOCK_THRESHOLD", "0.08")))
FIRST_BLOCK_DENSE_START = max(1, int(os.environ.get("H3_FIRST_BLOCK_DENSE_START", "3")))
FIRST_BLOCK_DENSE_END = max(1, int(os.environ.get("H3_FIRST_BLOCK_DENSE_END", "2")))
FORECAST_BLEND = min(1.0, max(0.0, float(os.environ.get("H3_FORECAST_BLEND", "0.65"))))
FUSED_ADALN = os.environ.get("H3_FUSED_ADALN", "0") == "1" and triton is not None
SOL_ATTN = os.environ.get("H3_SOL_ATTN", "1") == "1"
SOL_ATTN_BACKEND = os.environ.get("H3_SOL_ATTN_BACKEND", "triton").lower()
SOL_ATTN_TAU = float(os.environ.get("H3_SOL_ATTN_TAU", "1.0"))
SOL_ATTN_DENSE_STEPS = max(0, int(os.environ.get("H3_SOL_ATTN_DENSE_STEPS", "10")))
SOL_ATTN_DENSE_LAYERS = max(0, int(os.environ.get("H3_SOL_ATTN_DENSE_LAYERS", "2")))
SOL_ATTN_MIN_TOKENS = max(0, int(os.environ.get("H3_SOL_ATTN_MIN_TOKENS", "24576")))


if triton is not None:

    @triton.jit
    def _adaln_modulate_kernel(
        x, shift, scale, row_ids, elements: tl.constexpr, hidden: tl.constexpr, modulation_stride: tl.constexpr
    ):
        offsets = tl.program_id(0) * 256 + tl.arange(0, 256)
        mask = offsets < elements
        columns = offsets % hidden
        rows = offsets // hidden
        modulation_rows = tl.load(row_ids + rows, mask=mask, other=0)
        modulation_offsets = modulation_rows * modulation_stride + columns
        values = tl.load(x + offsets, mask=mask)
        shifts = tl.load(shift + modulation_offsets, mask=mask)
        scales = tl.load(scale + modulation_offsets, mask=mask)
        tl.store(x + offsets, values * (1.0 + scales) + shifts, mask=mask)

    @triton.jit
    def _adaln_gate_kernel(
        x, update, gate, row_ids, elements: tl.constexpr, hidden: tl.constexpr, modulation_stride: tl.constexpr
    ):
        offsets = tl.program_id(0) * 256 + tl.arange(0, 256)
        mask = offsets < elements
        columns = offsets % hidden
        rows = offsets // hidden
        modulation_rows = tl.load(row_ids + rows, mask=mask, other=0)
        gates = tl.load(gate + modulation_rows * modulation_stride + columns, mask=mask)
        values = tl.load(x + offsets, mask=mask)
        updates = tl.load(update + offsets, mask=mask)
        tl.store(x + offsets, values + updates * gates, mask=mask)


class H3StepCache:
    """ComfyUI EasyCache-style adaptive reuse of a complete H3 denoising result.

    This caches the model residual, not the generated video.  A request with a new prompt, seed, canvas or keyframe
    starts from an empty cache.  Decisions use a sparse sample of generated video latent rows, while the reused
    residual contains every video and audio row so their joint denoising trajectory stays coupled.
    """

    def __init__(self):
        self.total_steps = 0
        self.step = 0
        self.skipped = 0
        self.profile = "balanced"
        self.consecutive_skips = 0
        self.last_actual_step = None
        self.previous_input = None
        self.previous_output = None
        self.previous_output_norm = None
        self.relative_rate = None
        self.accumulated_change = None
        self.video_residual = None
        self.audio_residual = None
        self.video_residual_slope = None
        self.audio_residual_slope = None
        self.pending_input = None
        self.pending_input_change = None
        self.pending_track = False
        self.head_residual = None
        self.tail_residual = None
        self.first_block_output = None

    def begin(self, total_steps: int | None, profile: str = "balanced") -> None:
        self.__init__()
        self.total_steps = max(0, int(total_steps or 0))
        self.profile = str(profile or "balanced").lower()

    @property
    def enabled(self) -> bool:
        return self.profile != "exact" and self.total_steps > 2

    def _forecast(self, video_input, audio_input):
        distance = max(1, self.step - int(self.last_actual_step or 0))
        video_residual = self.video_residual
        audio_residual = self.audio_residual
        if self.video_residual_slope is not None:
            video_residual = video_residual + self.video_residual_slope * (distance * FORECAST_BLEND)
            audio_residual = audio_residual + self.audio_residual_slope * (distance * FORECAST_BLEND)
        self.skipped += 1
        self.consecutive_skips += 1
        self.step += 1
        return video_input + video_residual, audio_input + audio_residual

    def try_reuse(self, video_input, audio_input, condition_rows: int):
        self.pending_input = None
        self.pending_input_change = None
        self.pending_track = False
        if not self.enabled:
            return None

        # Balanced uses NVIDIA's H3 FirstBlockCache below, after block 0 has produced a high-signal residual.
        # Only the deliberately aggressive Ultra profile forecasts a whole transformer call before block 0.
        if not self.profile.startswith("ultra"):
            return None

        # Ultra Fast is deliberately bounded: no more than three forecasts can separate exact transformer calls, and
        # the high-noise warmup plus low-noise tail remain exact.  At the default 16 steps this executes 7 full DiT
        # evaluations instead of 16 while still sampling the original 16-step scheduler trajectory.
        if self.profile.startswith("ultra"):
            can_forecast = (
                self.step >= 3
                and self.step < self.total_steps - 2
                and self.consecutive_skips < 3
                and self.last_actual_step is not None
                and self.video_residual is not None
                and self.audio_residual is not None
                and self.video_residual.shape == video_input.shape
                and self.audio_residual.shape == audio_input.shape
            )
            if can_forecast:
                return self._forecast(video_input, audio_input)
            return None

        if EASYCACHE_THRESHOLD <= 0.0:
            return None

        end_step = math.floor(self.total_steps * EASYCACHE_END)
        if self.step >= end_step:
            return None

        # Condition latents are static. Excluding them makes the change estimate reflect the generated trajectory.
        sampled_input = video_input[0, condition_rows::EASYCACHE_SUBSAMPLE].detach().float()
        self.pending_input = sampled_input
        self.pending_track = True
        if self.previous_input is not None:
            self.pending_input_change = (sampled_input - self.previous_input).abs().mean()

        start_step = math.ceil(self.total_steps * EASYCACHE_START)
        can_reuse = (
            self.step >= start_step
            and self.pending_input_change is not None
            and self.relative_rate is not None
            and self.previous_output_norm is not None
            and self.video_residual is not None
            and self.audio_residual is not None
            and self.video_residual.shape == video_input.shape
            and self.audio_residual.shape == audio_input.shape
        )
        if not can_reuse:
            return None

        estimated_change = self.relative_rate * self.pending_input_change
        estimated_change = estimated_change / self.previous_output_norm.clamp_min(1e-6)
        accumulated = estimated_change if self.accumulated_change is None else self.accumulated_change + estimated_change
        if bool((accumulated < EASYCACHE_THRESHOLD).item()):
            self.accumulated_change = accumulated
            self.skipped += 1
            self.step += 1
            return video_input + self.video_residual, audio_input + self.audio_residual
        return None

    def first_block_decision(self, block_input: torch.Tensor, block_output: torch.Tensor) -> bool:
        """Return True when blocks 1..49 can reuse their previous joint residual.

        This is the single-GPU equivalent of NVIDIA Sol-Engine's H3 FirstBlockCache at threshold 0.08. The first
        block is always evaluated. Its normalized residual change is a much stronger predictor than raw latent
        motion, while the cached tail residual still covers the complete text/video/audio packed sequence.
        """
        if not self.enabled or self.profile.startswith("ultra"):
            return False
        if FIRST_BLOCK_THRESHOLD <= 0.0:
            self.first_block_output = block_output.detach().clone()
            return False

        keep_dense = self.step < FIRST_BLOCK_DENSE_START or self.step >= self.total_steps - FIRST_BLOCK_DENSE_END
        residual = block_output - block_input
        reusable = (
            not keep_dense
            and self.head_residual is not None
            and self.tail_residual is not None
            and self.tail_residual.shape == block_output.shape
        )
        should_reuse = False
        if reusable:
            difference = (residual - self.head_residual).abs().mean()
            reference = self.head_residual.abs().mean().clamp_min(1e-8)
            should_reuse = bool(((difference / reference) <= FIRST_BLOCK_THRESHOLD).item())

        if should_reuse:
            self.skipped += 1
            self.consecutive_skips += 1
            self.step += 1
            return True

        # This engine's residual/gate operations update `packed` in place. Preserve the head output before later
        # blocks mutate the same storage; diffusers' reference blocks are out-of-place and do not need this clone.
        self.first_block_output = block_output.detach().clone()
        self.head_residual = residual.detach()
        return False

    def update_first_block_tail(self, final_block_output: torch.Tensor) -> None:
        if self.first_block_output is None:
            return
        self.tail_residual = (final_block_output - self.first_block_output).detach()
        self.last_actual_step = self.step
        self.consecutive_skips = 0
        self.step += 1
        self.first_block_output = None

    def update(self, video_input, audio_input, video_output, audio_output, condition_rows: int) -> None:
        # Balanced's clock and state are updated at the block-stack boundary by FirstBlockCache.
        if self.enabled and not self.profile.startswith("ultra"):
            return
        if self.pending_track:
            sampled_output = video_output[0, condition_rows::EASYCACHE_SUBSAMPLE].detach().float()
            if self.previous_output is not None and self.pending_input_change is not None:
                output_change = (sampled_output - self.previous_output).abs().mean()
                self.relative_rate = output_change / self.pending_input_change.clamp_min(1e-6)
            self.previous_input = self.pending_input.clone()
            self.previous_output = sampled_output.clone()
            self.previous_output_norm = sampled_output.abs().mean()
            if not self.profile.startswith("ultra"):
                self.video_residual = (video_output - video_input).detach()
                self.audio_residual = (audio_output - audio_input).detach()
            self.accumulated_change = None
        if self.profile.startswith("ultra"):
            new_video_residual = (video_output - video_input).detach()
            new_audio_residual = (audio_output - audio_input).detach()
            if self.video_residual is not None and self.last_actual_step is not None:
                gap = max(1, self.step - self.last_actual_step)
                self.video_residual_slope = (new_video_residual - self.video_residual) / gap
                self.audio_residual_slope = (new_audio_residual - self.audio_residual) / gap
            self.video_residual = new_video_residual
            self.audio_residual = new_audio_residual
        self.last_actual_step = self.step
        self.consecutive_skips = 0
        self.step += 1
        self.pending_input = None
        self.pending_input_change = None
        self.pending_track = False

    def finish(self) -> dict:
        stats = {
            "steps": self.step,
            "computed": max(0, self.step - self.skipped),
            "forecasted": self.skipped,
            "profile": self.profile,
        }
        if self.enabled and self.step:
            computed = max(1, self.step - self.skipped)
            print(
                f"[h3-nvfp4] adaptive step cache skipped {self.skipped}/{self.step} transformer evaluations "
                f"({self.step / computed:.2f}x denoiser-work reduction)",
                flush=True,
            )
        self.begin(None)
        return stats


class H3SolAttention:
    """NVIDIA Sol-Attn policy adapted to H3's single-GPU packed attention.

    The packed prefix (text, conditioning video and generated audio) remains an exact KV sink and its query rows are
    recomputed densely. Only target-video query/key interactions become sparse, after ten dense denoising steps and
    outside the first two transformer blocks. Any unavailable/JIT-failing backend falls back to cuDNN for the request.
    """

    def __init__(self):
        self.enabled = SOL_ATTN
        self.step = 0
        self.video_start = 0
        self.sparse_calls = 0
        self.dense_calls = 0
        self.failure = None

    def begin(self):
        self.step = 0
        self.video_start = 0
        self.sparse_calls = 0
        self.dense_calls = 0
        self.failure = None

    def observe(self, video_indices: torch.Tensor, sequence: int, step: int) -> None:
        self.step = int(step)
        if not self.video_start:
            deltas = video_indices[1:] - video_indices[:-1]
            breaks = (deltas != 1).nonzero().flatten()
            start = int(breaks[-1]) + 1 if len(breaks) else 0
            self.video_start = int(video_indices[start]) if video_indices.numel() else sequence

    def __call__(self, query, key, value, layer: int):
        tokens = int(query.shape[1])
        if (
            not self.enabled
            or self.failure is not None
            or self.step < SOL_ATTN_DENSE_STEPS
            or layer < SOL_ATTN_DENSE_LAYERS
            or tokens < SOL_ATTN_MIN_TOKENS
            or not 0 < self.video_start < tokens
        ):
            self.dense_calls += 1
            return None
        try:
            if SOL_ATTN_BACKEND == "triton":
                from sol_attn.triton_ref import sol_attn
            else:
                from sol_attn import sol_attn

            q, k, v = (tensor.contiguous() for tensor in (query, key, value))
            kwargs = {
                "tau": SOL_ATTN_TAU,
                "thresh_type": "diag",
                "sink_start": 0,
                "sink_tokens": self.video_start,
            }
            if SOL_ATTN_BACKEND != "triton":
                kwargs["kv_splits"] = 1
            attended = sol_attn(q, k, v, **kwargs)
            # An exact KV sink does not make the prefix's own queries dense. H3 jointly generates audio in that
            # prefix, so reproduce those rows with exact attention as NVIDIA's H3 integration does.
            prefix = self.video_start
            dense_prefix = F.scaled_dot_product_attention(
                q[:, :prefix].transpose(1, 2),
                k.transpose(1, 2),
                v.transpose(1, 2),
                dropout_p=0.0,
                is_causal=False,
            ).transpose(1, 2)
            attended[:, :prefix] = dense_prefix
            self.sparse_calls += 1
            return attended
        except Exception as error:
            self.failure = f"{type(error).__name__}: {error}"
            print(f"[h3-sol-attn] falling back to dense attention: {self.failure}", flush=True)
            self.dense_calls += 1
            return None


def _quant_config(handle, prefix: str) -> dict | None:
    key = f"{prefix}.comfy_quant"
    if key not in handle.keys():
        return None
    return json.loads(handle.get_tensor(key).numpy().tobytes())


class H3Linear(nn.Module):
    """A plain or comfy-kitchen NVFP4 linear, selected by checkpoint metadata."""

    def __init__(
        self,
        in_features: int,
        out_features: int,
        bias: bool = False,
        compute_dtype: torch.dtype | None = None,
    ):
        super().__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.compute_dtype = compute_dtype
        self.register_parameter("weight", None)
        self.register_parameter("bias", None)
        self.register_buffer("input_scale", None)
        self.register_buffer("pre_quant_scale", None)
        self.quantized = False
        self.full_precision_mm = False

    def load(self, handle, prefix: str) -> None:
        config = _quant_config(handle, prefix)
        weight = handle.get_tensor(f"{prefix}.weight")

        if config is None:
            self.weight = nn.Parameter(
                weight if self.compute_dtype is None else weight.to(self.compute_dtype), requires_grad=False
            )
        elif config.get("format") == "nvfp4":
            block_scale = handle.get_tensor(f"{prefix}.weight_scale")
            if block_scale.dtype == torch.uint8:
                block_scale = block_scale.view(torch.float8_e4m3fn)
            tensor_scale = handle.get_tensor(f"{prefix}.weight_scale_2").float()
            params = TensorCoreNVFP4Layout.Params(
                scale=tensor_scale,
                block_scale=block_scale,
                orig_dtype=torch.bfloat16,
                orig_shape=(self.out_features, self.in_features),
            )
            quantized = QuantizedTensor(weight.to(torch.uint8), "TensorCoreNVFP4Layout", params)
            self.weight = nn.Parameter(quantized, requires_grad=False)
            self.quantized = True
            self.full_precision_mm = bool(config.get("full_precision_matrix_mult", False))
            for name in ("input_scale", "pre_quant_scale"):
                key = f"{prefix}.{name}"
                if key in handle.keys():
                    setattr(self, name, handle.get_tensor(key))
        else:
            raise ValueError(f"Unsupported quantization on {prefix}: {config}")

        bias_key = f"{prefix}.bias"
        if bias_key in handle.keys():
            bias = handle.get_tensor(bias_key)
            self.bias = nn.Parameter(
                bias if self.compute_dtype is None else bias.to(self.compute_dtype), requires_grad=False
            )

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        if self.pre_quant_scale is not None:
            hidden_states = hidden_states * self.pre_quant_scale.to(
                device=hidden_states.device, dtype=hidden_states.dtype
            )
        if not self.quantized:
            hidden_states = hidden_states.to(self.weight.dtype)
            return F.linear(
                hidden_states,
                self.weight,
                self.bias,
            )

        if self.full_precision_mm:
            # Some AWQ checkpoints use NVFP4 as a compact weight format but deliberately retain BF16 activations and
            # GEMMs. Dequantization is layer-local, so residency stays compact without adding activation error.
            weight = self.weight.dequantize().to(hidden_states.dtype)
            return F.linear(hidden_states, weight, None if self.bias is None else self.bias.to(hidden_states.dtype))

        shape = hidden_states.shape
        flat = hidden_states.reshape(-1, shape[-1])
        scale = None if self.input_scale is None else self.input_scale.to(flat.device)
        quantized_input = QuantizedTensor.from_float(flat, "TensorCoreNVFP4Layout", scale=scale)
        output = F.linear(
            quantized_input,
            self.weight,
            None if self.bias is None else self.bias.to(hidden_states.dtype),
        )
        return output.reshape(*shape[:-1], self.out_features)


class H3RMSNorm(nn.Module):
    def __init__(self, width: int, eps: float = EPS):
        super().__init__()
        self.width = width
        self.eps = eps
        self.register_parameter("weight", None)

    def load(self, handle, prefix: str) -> None:
        self.weight = nn.Parameter(handle.get_tensor(f"{prefix}.weight"), requires_grad=False)

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        return F.rms_norm(
            hidden_states,
            (self.width,),
            self.weight,
            self.eps,
        )


class H3Attention(nn.Module):
    def __init__(self):
        super().__init__()
        self.qkv_proj = H3Linear(HIDDEN, 3 * HEADS * HEAD_DIM)
        self.q_norm = H3RMSNorm(HEAD_DIM)
        self.k_norm = H3RMSNorm(HEAD_DIM)
        self.out_proj = H3Linear(HEADS * HEAD_DIM, HIDDEN)

    def load(self, handle, prefix: str) -> None:
        self.qkv_proj.load(handle, f"{prefix}.qkv_proj")
        self.q_norm.load(handle, f"{prefix}.q_norm")
        self.k_norm.load(handle, f"{prefix}.k_norm")
        self.out_proj.load(handle, f"{prefix}.out_proj")

    def forward(self, hidden_states, rope_table, backend: str, sparse=None, layer: int = -1):
        sequence = hidden_states.shape[0]
        qkv = self.qkv_proj(hidden_states)
        query, key, value = qkv.split(HEADS * HEAD_DIM, dim=-1)
        query = query.view(1, sequence, HEADS, HEAD_DIM)
        key = key.view(1, sequence, HEADS, HEAD_DIM)
        value = value.view(1, sequence, HEADS, HEAD_DIM)

        # One in-place kernel replaces Q RMSNorm, K RMSNorm and both partial RoPE applications.
        kitchen.rms_rope_split_half_(
            query,
            key,
            rope_table,
            self.q_norm.weight,
            self.k_norm.weight,
            epsilon=self.q_norm.eps,
            rot_dim=rope_table.shape[-3] * 2,
        )
        attended = sparse(query, key, value, layer) if sparse is not None else None
        if attended is None:
            attended = dispatch_attention_fn(
                query,
                key,
                value,
                attn_mask=None,
                dropout_p=0.0,
                is_causal=False,
                backend=backend,
            )
        return self.out_proj(attended.reshape(sequence, HEADS * HEAD_DIM))


class H3MLP(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = H3Linear(HIDDEN, 2 * FFN)
        self.fc2 = H3Linear(FFN, HIDDEN)

    def load(self, handle, prefix: str) -> None:
        self.fc1.load(handle, f"{prefix}.fc1")
        self.fc2.load(handle, f"{prefix}.fc2")

    def forward(self, hidden_states):
        gate, up = self.fc1(hidden_states).chunk(2, dim=-1)
        return self.fc2(F.silu(gate).mul_(up))


class H3RefinerBlock(nn.Module):
    def __init__(self):
        super().__init__()
        self.norm1 = H3RMSNorm(HIDDEN)
        self.attn = H3Attention()
        self.norm2 = H3RMSNorm(HIDDEN)
        self.mlp = H3MLP()

    def load(self, handle, prefix: str) -> None:
        self.norm1.load(handle, f"{prefix}.norm1")
        self.attn.load(handle, f"{prefix}.attn")
        self.norm2.load(handle, f"{prefix}.norm2")
        self.mlp.load(handle, f"{prefix}.mlp")


class H3AdaLN(nn.Module):
    def __init__(self, expand: int, modalities: int):
        super().__init__()
        self.expand = expand
        self.modalities = modalities
        # Curve checkpoints deliberately evaluate interpolation and modulation projection in FP32. Expanding the
        # checkpoint's tiny FP16 [*, 8] matrices once at load avoids 51 request-step casts.
        self.linear = H3Linear(
            TIME_DIM, expand * HIDDEN * modalities, bias=True, compute_dtype=torch.float32
        )

    def load(self, handle, prefix: str) -> None:
        self.linear.load(handle, f"{prefix}.linear")

    def forward(self, time_embedding, output_dtype=None):
        projected = self.linear(time_embedding)
        if output_dtype is not None:
            # One contiguous conversion is numerically identical to converting the six chunk views independently,
            # and removes five CUDA launches from every one of the 50 blocks.
            projected = projected.to(output_dtype)
        projected = projected.view(-1, self.expand * HIDDEN)
        return projected.chunk(self.expand, dim=-1)


class H3Block(nn.Module):
    def __init__(self):
        super().__init__()
        self.norm1 = H3RMSNorm(HIDDEN)
        self.attn = H3Attention()
        self.norm2 = H3RMSNorm(HIDDEN)
        self.mlp = H3MLP()
        self.adaln_proj = H3AdaLN(6, 3)

    def load(self, handle, prefix: str) -> None:
        self.norm1.load(handle, f"{prefix}.norm1")
        self.attn.load(handle, f"{prefix}.attn")
        self.norm2.load(handle, f"{prefix}.norm2")
        self.mlp.load(handle, f"{prefix}.mlp")
        self.adaln_proj.load(handle, f"{prefix}.adaln_proj")


class H3FinalLayer(nn.Module):
    def __init__(self):
        super().__init__()
        self.norm = H3RMSNorm(HIDDEN)
        self.adaln_proj = H3AdaLN(2, 1)
        self.video_out = H3Linear(HIDDEN, VIDEO_DIM, bias=True, compute_dtype=torch.float32)
        self.audio_out = H3Linear(HIDDEN, AUDIO_DIM, bias=True, compute_dtype=torch.float32)

    def load(self, handle, prefix: str) -> None:
        self.norm.load(handle, f"{prefix}.norm")
        self.adaln_proj.load(handle, f"{prefix}.adaln_proj")
        self.video_out.load(handle, f"{prefix}.video_out")
        self.audio_out.load(handle, f"{prefix}.audio_out")


class H3NVFP4Transformer(nn.Module):
    """Diffusers-compatible H3 transformer backed by fused comfy-kitchen NVFP4 kernels."""

    def __init__(self):
        super().__init__()
        # The modular pipeline reads these values through the diffusers component config rather than inspecting the
        # module itself. Keep the public transformer contract even though this lean adapter is not a ConfigMixin.
        self.config = SimpleNamespace(
            patch_size=(1, 2, 2),
            in_channels=24,
            audio_in_channels=AUDIO_DIM,
            text_dim=TEXT_DIM,
        )
        self.video_patch_proj = H3Linear(VIDEO_DIM, HIDDEN, bias=True, compute_dtype=torch.float32)
        self.audio_patch_proj = H3Linear(AUDIO_DIM, HIDDEN, bias=True, compute_dtype=torch.float32)
        self.condition_proj = H3Linear(TEXT_DIM, HIDDEN, bias=True)
        self.token_refiner = nn.ModuleList([H3RefinerBlock() for _ in range(REFINER_LAYERS)])
        self.token_refiner_norm = H3RMSNorm(HIDDEN)
        self.blocks = nn.ModuleList([H3Block() for _ in range(LAYERS)])
        self.final_layer = H3FinalLayer()
        self.register_buffer("adaln_t_table", None)
        self.register_buffer("rope_inv_freq", None)
        self.attention_backend = "_native_cudnn"
        self._text_cache = None
        self._rope_cache = None
        self._segment_cache = None
        self._condition_video_rows = None
        self._condition_video_embedding = None
        self._output_indices = None
        self._generated_rows = None
        self._step_cache = H3StepCache()
        self._sol_attention = H3SolAttention()

    @property
    def dtype(self) -> torch.dtype:
        """Match ModelMixin's placement contract used by ModularPipeline.to()."""
        return self.condition_proj.weight.dtype

    @property
    def device(self) -> torch.device:
        return self.adaln_t_table.device

    def load(self, path: str) -> None:
        from safetensors import safe_open

        with safe_open(path, framework="pt", device="cpu") as handle:
            self.video_patch_proj.load(handle, "video_patch_proj")
            self.audio_patch_proj.load(handle, "audio_patch_proj")
            self.condition_proj.load(handle, "condition_proj")
            for index, block in enumerate(self.token_refiner):
                block.load(handle, f"token_refiner.blocks.{index}")
            self.token_refiner_norm.load(handle, "token_refiner.final_norm")
            for index, block in enumerate(self.blocks):
                block.load(handle, f"blocks.{index}")
            self.final_layer.load(handle, "final_layer")
            self.adaln_t_table = handle.get_tensor("adaln_t_table")
            self.rope_inv_freq = handle.get_tensor("rope.inv_freq")
        # Every loaded tensor is already a frozen Parameter (or a buffer). Avoid mutating the quantized tensor
        # subclass through a redundant requires_grad_ dispatch.
        self.eval()

    def set_attention_backend(self, backend: str) -> None:
        self.attention_backend = backend

    def begin_request(self, total_steps: int | None = None, profile: str = "balanced") -> None:
        self._text_cache = None
        self._rope_cache = None
        self._segment_cache = None
        self._condition_video_rows = None
        self._condition_video_embedding = None
        self._output_indices = None
        self._generated_rows = None
        self._step_cache.begin(total_steps, profile)
        self._sol_attention.begin()

    def end_request(self) -> dict:
        stats = self._step_cache.finish()
        stats["sol_sparse_calls"] = self._sol_attention.sparse_calls
        stats["sol_dense_calls"] = self._sol_attention.dense_calls
        stats["sol_failure"] = self._sol_attention.failure
        self._text_cache = None
        self._rope_cache = None
        self._segment_cache = None
        self._condition_video_rows = None
        self._condition_video_embedding = None
        self._output_indices = None
        self._generated_rows = None
        return stats

    def _refine_text(self, text_states: torch.Tensor) -> torch.Tensor:
        key = (text_states.data_ptr(), tuple(text_states.shape), text_states.device)
        if self._text_cache is not None and self._text_cache[0] == key:
            return self._text_cache[1]
        hidden = self.condition_proj(text_states)
        # Text is tiny compared with the video sequence; use the same fused QKV path with an identity RoPE omitted.
        for block in self.token_refiner:
            residual = hidden
            normalized = block.norm1(hidden)
            qkv = block.attn.qkv_proj(normalized)
            query, key_states, value = qkv.split(HEADS * HEAD_DIM, dim=-1)
            query = block.attn.q_norm(query.view(1, -1, HEADS, HEAD_DIM))
            key_states = block.attn.k_norm(key_states.view(1, -1, HEADS, HEAD_DIM))
            value = value.view(1, -1, HEADS, HEAD_DIM)
            attended = dispatch_attention_fn(
                query,
                key_states,
                value,
                attn_mask=None,
                dropout_p=0.0,
                is_causal=False,
                backend=self.attention_backend,
            ).reshape(-1, HEADS * HEAD_DIM)
            hidden = residual + block.attn.out_proj(attended)
            hidden = hidden + block.mlp(block.norm2(hidden))
        hidden = self.token_refiner_norm(hidden)
        self._text_cache = (key, hidden)
        return hidden

    def _rope(self, position_ids: torch.Tensor, dtype: torch.dtype) -> torch.Tensor:
        key = (position_ids.data_ptr(), tuple(position_ids.shape), position_ids.device, dtype)
        if self._rope_cache is not None and self._rope_cache[0] == key:
            return self._rope_cache[1]
        positions = position_ids.to(torch.float32)
        frequencies = positions.unsqueeze(-1) * self.rope_inv_freq.to(position_ids.device).view(1, 1, -1)
        temporal, height, width = frequencies.unbind(dim=1)
        angles = torch.cat((temporal, height, width), dim=-1)
        cosine, sine = angles.cos(), angles.sin()
        table = torch.stack((cosine, -sine, sine, cosine), dim=-1)
        table = table.reshape(1, position_ids.shape[0], 1, angles.shape[-1], 2, 2).to(dtype)
        self._rope_cache = (key, table)
        return table

    def _time_embedding(self, timestep: torch.Tensor) -> torch.Tensor:
        table = self.adaln_t_table.to(timestep.device)
        position = timestep.float().clamp(0.0, 1.0) * (table.shape[0] - 1)
        lower = position.floor().long().clamp(max=table.shape[0] - 2)
        return torch.lerp(table[lower], table[lower + 1], (position - lower).unsqueeze(1))

    def _segments(self, indices: torch.Tensor):
        if self._segment_cache is None:
            host = indices.detach().cpu()
            changes = (host[1:] != host[:-1]).nonzero().flatten().add(1).tolist()
            bounds = [0, *changes, len(host)]
            # Python row ids avoid indexing modulation tensors with CUDA scalar tensors in every block.
            self._segment_cache = [
                (start, stop, int(host[start])) for start, stop in zip(bounds[:-1], bounds[1:])
            ]
        return self._segment_cache

    def _video_layout(self, video_indices: torch.Tensor) -> int:
        """Number of leading, static keyframe-patch rows in the video latent tensor."""
        if self._condition_video_rows is None:
            host = video_indices.detach().cpu()
            discontinuities = (host[1:] - host[:-1] != 1).nonzero().flatten()
            self._condition_video_rows = int(discontinuities[0]) + 1 if len(discontinuities) else 0
        return self._condition_video_rows

    def _project_video(self, hidden_states: torch.Tensor, condition_rows: int, dtype: torch.dtype) -> torch.Tensor:
        source = hidden_states[0]
        if condition_rows == 0:
            return self.video_patch_proj(source.float()).to(dtype)
        if self._condition_video_embedding is None:
            self._condition_video_embedding = self.video_patch_proj(source[:condition_rows].float()).to(dtype)
        generated = self.video_patch_proj(source[condition_rows:].float()).to(dtype)
        return torch.cat((self._condition_video_embedding, generated), dim=0)

    @staticmethod
    def _modulate(hidden, shift, scale, row_ids, segments):
        if FUSED_ADALN and hidden.is_cuda and hidden.is_contiguous():
            _adaln_modulate_kernel[(triton.cdiv(hidden.numel(), 256),)](
                hidden, shift, scale, row_ids, hidden.numel(), HIDDEN, shift.stride(0), num_warps=4
            )
            return hidden
        for start, stop, row in segments:
            hidden[start:stop].mul_(1.0 + scale[row]).add_(shift[row])
        return hidden

    @staticmethod
    def _gate(hidden, update, gate, row_ids, segments):
        if FUSED_ADALN and hidden.is_cuda and hidden.is_contiguous() and update.is_contiguous():
            _adaln_gate_kernel[(triton.cdiv(hidden.numel(), 256),)](
                hidden, update, gate, row_ids, hidden.numel(), HIDDEN, gate.stride(0), num_warps=4
            )
            return hidden
        for start, stop, row in segments:
            hidden[start:stop].addcmul_(update[start:stop], gate[row])
        return hidden

    def forward(
        self,
        hidden_states,
        audio_hidden_states,
        encoder_hidden_states,
        timestep,
        timestep_indices,
        token_tags,
        position_ids,
        video_indices,
        audio_indices,
        text_indices,
        attention_kwargs=None,
        return_dict=True,
    ):
        from diffusers.models.transformers.transformer_minimax_h3 import MiniMaxH3TransformerOutput

        if hidden_states.shape[0] != 1:
            raise ValueError("The NVFP4 MiniMax-H3 engine supports batch size 1.")

        condition_rows = self._video_layout(video_indices)
        reused = self._step_cache.try_reuse(hidden_states, audio_hidden_states, condition_rows)
        if reused is not None:
            video_output, audio_output = reused
            if not return_dict:
                return video_output, audio_output
            return MiniMaxH3TransformerOutput(sample=video_output, audio_sample=audio_output)

        text = self._refine_text(encoder_hidden_states[0].to(torch.bfloat16))
        video = self._project_video(hidden_states, condition_rows, text.dtype)
        audio = self.audio_patch_proj(audio_hidden_states[0].float()).to(text.dtype)
        # Text, video and audio indices partition the packed sequence, so initialization would only add a full HBM
        # write before the three index copies overwrite every row.
        packed = text.new_empty((position_ids.shape[0], HIDDEN))
        packed.index_copy_(0, text_indices, text)
        packed.index_copy_(0, video_indices, video)
        packed.index_copy_(0, audio_indices, audio)

        time_embedding = self._time_embedding(timestep)
        adaln_indices = timestep_indices * 3 + token_tags.clamp(min=0)
        segments = self._segments(adaln_indices)
        rope = self._rope(position_ids, packed.dtype)
        use_sol_attention = self._step_cache.profile != "exact" and self._sol_attention.enabled
        self._sol_attention.observe(video_indices, packed.shape[0], self._step_cache.step)

        reused_tail = False
        for layer, block in enumerate(self.blocks):
            if layer == 0:
                # Block 0 writes its residual updates in place, so retain the pre-block value for the official FBC
                # signal `(head_output - head_input)`.
                block_input = packed.detach().clone()
            # One conversion per small modulation table, rather than one conversion per sequence segment.
            modulations = block.adaln_proj(time_embedding, packed.dtype)
            shift_attn, scale_attn, gate_attn, shift_mlp, scale_mlp, gate_mlp = modulations
            normalized = self._modulate(block.norm1(packed), shift_attn, scale_attn, adaln_indices, segments)
            packed = self._gate(
                packed,
                block.attn(
                    normalized,
                    rope,
                    self.attention_backend,
                    self._sol_attention if use_sol_attention else None,
                    layer,
                ),
                gate_attn,
                adaln_indices,
                segments,
            )
            normalized = self._modulate(block.norm2(packed), shift_mlp, scale_mlp, adaln_indices, segments)
            packed = self._gate(packed, block.mlp(normalized), gate_mlp, adaln_indices, segments)

            if layer == 0:
                if self._step_cache.first_block_decision(block_input, packed):
                    packed = packed + self._step_cache.tail_residual
                    reused_tail = True
                    break
            if layer == len(self.blocks) - 1 and not reused_tail:
                self._step_cache.update_first_block_tail(packed)

        shift, scale = self.final_layer.adaln_proj(time_embedding)

        # Keyframe output rows are discarded by the scheduler. Avoid their FP32 output projection and put zeros in
        # those unused slots to retain the pipeline's expected tensor shape.
        generated_video_indices = video_indices[condition_rows:]
        if self._output_indices is None:
            self._generated_rows = generated_video_indices.shape[0]
            self._output_indices = torch.cat((generated_video_indices, audio_indices))
        generated_rows = self._generated_rows
        normalized_output = self.final_layer.norm(packed.index_select(0, self._output_indices))
        video_times = timestep_indices.index_select(0, generated_video_indices)
        video_hidden = normalized_output[:generated_rows]
        video_hidden = video_hidden * (1.0 + scale.index_select(0, video_times)) + shift.index_select(0, video_times)
        generated_video_output = self.final_layer.video_out(video_hidden.float())
        if condition_rows:
            video_output = generated_video_output.new_zeros((1, hidden_states.shape[1], VIDEO_DIM))
            video_output[0, condition_rows:] = generated_video_output
        else:
            video_output = generated_video_output.unsqueeze(0)

        audio_times = timestep_indices.index_select(0, audio_indices)
        audio_hidden = normalized_output[generated_rows:]
        audio_hidden = audio_hidden * (1.0 + scale.index_select(0, audio_times)) + shift.index_select(0, audio_times)
        audio_output = self.final_layer.audio_out(audio_hidden.float()).unsqueeze(0)

        self._step_cache.update(
            hidden_states,
            audio_hidden_states,
            video_output,
            audio_output,
            condition_rows,
        )

        if not return_dict:
            return video_output, audio_output
        return MiniMaxH3TransformerOutput(sample=video_output, audio_sample=audio_output)


def load_transformer() -> H3NVFP4Transformer:
    if torch.version.cuda is None or int(torch.version.cuda.split(".")[0]) < 13:
        raise RuntimeError("NVFP4 requires the CUDA 13 PyTorch build.")
    from huggingface_hub import hf_hub_download

    path = hf_hub_download(repo_id=NVFP4_REPO, filename=NVFP4_FILE)
    transformer = H3NVFP4Transformer()
    transformer.load(path)
    print(f"[h3-nvfp4] loaded {NVFP4_REPO}/{NVFP4_FILE}", flush=True)
    return transformer


def status() -> str:
    return (
        f"NVFP4 · linear residual forecast {FORECAST_BLEND:g} / adaptive cache {EASYCACHE_THRESHOLD:g} · "
        f"pruned AdaLN curve · fused QKV/QK-norm/RoPE · `{NVFP4_REPO}`"
    )