"""Direct, decoder-only MiniMax H3 video VAE implementation."""

from __future__ import annotations

import math
import os
from pathlib import Path

import torch
from safetensors import safe_open
from torch import nn
from torch.nn import functional as F


IMAGENET_MEAN = (0.485, 0.456, 0.406)
IMAGENET_STD = (0.229, 0.224, 0.225)
LATENTS_MEAN = (0.858090341091156, -0.9606591463088989, 1.0661640167236328, -0.5090325474739075, -0.2727581858634949, -1.3675414323806763, -0.2553254961967468, -0.26907554268836975, -0.5376840829849243, -0.0464097298681736, 0.6657370328903198, 0.19690127670764923, -0.5460608005523682, -0.4035342037677765, -0.23683024942874908, 0.25928452610969543, -0.30133944749832153, 0.211341992020607, -1.1206848621368408, 0.3581933379173279, -0.04225143790245056, 0.2604829967021942, 0.22864092886447906, 0.7056031823158264)
LATENTS_STD = (1.2223774194717407, 1.2767263650894165, 1.68317747116088865, 1.7549455165863037, 1.5636216402053833, 2.194143533706665, 0.96531379222869875, 1.05698859691619875, 0.841948926448822, 0.7729952931404114, 1.8955937623977661, 0.946841835975647, 0.7996809482574463, 0.44988900423049925, 0.7197399735450745, 0.69362932443618775, 2.961095094680786, 2.7694199085235595, 3.0496184825897215, 2.1088054180145265, 3.276226282119751, 3.1627357006073, 2.28168129920959475, 2.6127843856811525)


def _rms_norm(x: torch.Tensor, weight: torch.Tensor | None, eps: float) -> torch.Tensor:
    result = x * torch.rsqrt(x.float().square().mean(dim=-1, keepdim=True) + eps).to(x.dtype)
    return result if weight is None else result * weight.to(dtype=x.dtype)


def create_token_ids(patch_dims: tuple[int, int, int], device: torch.device, dtype: torch.dtype) -> torch.Tensor:
    coords = [2.0 * (torch.arange(0.5, size, dtype=dtype, device=device) / size) - 1.0 for size in patch_dims]
    return torch.stack(torch.meshgrid(*coords, indexing="ij"), dim=-1).flatten(0, 2).unsqueeze(0)


class RotaryEmbeddingND(nn.Module):
    def __init__(self, dim: int, rotary_base: float = 100.0, n_dim: int = 3, *, device=None):
        super().__init__()
        self.rotary_base = rotary_base
        self.step = 2 * n_dim / dim
        inv_freq = 1 / rotary_base ** torch.arange(0, 1, self.step, dtype=torch.float32, device=device)
        self.register_buffer("inv_freq", inv_freq, persistent=False)
        self.angle_scale = 2.0 * math.pi

    def forward(self, img_ids: torch.Tensor) -> torch.Tensor:
        inv_freq = self.inv_freq
        if inv_freq.device.type == "meta":
            inv_freq = 1 / self.rotary_base ** torch.arange(0, 1, self.step, dtype=torch.float32, device=img_ids.device)
        else:
            inv_freq = inv_freq.to(img_ids.device)
        angles = self.angle_scale * img_ids[:, :, :, None].float() * inv_freq[None, None, None, :]
        angles = angles.flatten(2, 3)
        cos, sin = torch.cos(angles), torch.sin(angles)
        return torch.stack((cos, -sin, sin, cos), dim=-1).reshape(*angles.shape[:2], 1, -1, 2, 2).to(img_ids.dtype)


class RMSNorm(nn.Module):
    def __init__(self, dim: int, eps: float, affine: bool, *, device=None):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.empty(dim, device=device)) if affine else None

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return _rms_norm(x, self.weight, self.eps)


class FeedForward(nn.Module):
    def __init__(self, dim: int, bias: bool = True, *, device=None):
        super().__init__()
        self.w1 = nn.Linear(dim, dim * 8, bias=bias, device=device)
        self.w2 = nn.Linear(dim * 4, dim, bias=bias, device=device)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        gate, value = self.w1(x).chunk(2, dim=-1)
        return self.w2(F.silu(gate).mul_(value))


def _apply_rope_split_half(x: torch.Tensor, table: torch.Tensor) -> torch.Tensor:
    """Apply the reference split-half RoPE layout to leading rotary channels."""
    pairs = table.shape[-3]
    rot = pairs * 2
    first, second = x[..., :pairs], x[..., pairs:rot]
    cos, neg_sin, sin = table[..., 0, 0], table[..., 0, 1], table[..., 1, 0]
    rotated = torch.cat((first * cos + second * neg_sin, first * sin + second * cos), dim=-1)
    return torch.cat((rotated, x[..., rot:]), dim=-1)


class Attention(nn.Module):
    def __init__(self, heads: int, dim_head: int, bias: bool = True, eps: float = 1e-5, *, device=None):
        super().__init__()
        dim = heads * dim_head
        self.heads, self.dim_head = heads, dim_head
        self.norm_q = RMSNorm(dim_head, eps, False, device=device)
        self.norm_k = RMSNorm(dim_head, eps, False, device=device)
        self.to_qkv = nn.Linear(dim, dim * 3, bias=bias, device=device)
        self.to_out = nn.Linear(dim, dim, bias=bias, device=device)

    def forward(self, x: torch.Tensor, rotary_pos_emb: torch.Tensor) -> torch.Tensor:
        batch, sequence, _ = x.shape
        qkv = self.to_qkv(x).view(batch, sequence, self.heads, 3 * self.dim_head)
        query, key, value = qkv.chunk(3, dim=-1)
        query, key = self.norm_q(query), self.norm_k(key)
        query, key = _apply_rope_split_half(query, rotary_pos_emb), _apply_rope_split_half(key, rotary_pos_emb)
        query, key, value = query.transpose(1, 2), key.transpose(1, 2), value.transpose(1, 2)
        try:
            from comfy.ldm.modules.attention import optimized_attention

            output = optimized_attention(query, key, value, self.heads, skip_reshape=True)
        except Exception:
            try:
                from sageattention import sageattn

                output = sageattn(query, key, value, is_causal=False, tensor_layout="HND", smooth_k=False).transpose(1, 2)
            except Exception:
                output = F.scaled_dot_product_attention(query, key, value).transpose(1, 2)
        return self.to_out(output.reshape(batch, sequence, -1).nan_to_num_(0.0))


class TransformerBlock(nn.Module):
    def __init__(self, heads: int, dim_head: int, bias: bool = True, eps: float = 1e-5, *, device=None):
        super().__init__()
        dim = heads * dim_head
        self.norm1 = RMSNorm(dim, eps, True, device=device)
        self.attn = Attention(heads, dim_head, bias, eps, device=device)
        self.scale1 = nn.Parameter(torch.empty(dim, device=device))
        self.norm2 = RMSNorm(dim, eps, True, device=device)
        self.ff = FeedForward(dim, bias, device=device)
        self.scale2 = nn.Parameter(torch.empty(dim, device=device))

    def forward(self, x: torch.Tensor, rotary_pos_emb: torch.Tensor) -> torch.Tensor:
        x = x.addcmul_(self.attn(self.norm1(x), rotary_pos_emb), self.scale1.to(x.dtype))
        return x.addcmul_(self.ff(self.norm2(x)), self.scale2.to(x.dtype))


class ViT3DDecoder(nn.Module):
    def __init__(self, patch_size: int = 16, patch_size_t: int = 4, in_channels: int = 24, out_channels: int = 3, num_layers: int = 36, heads: int = 32, dim_head: int = 64, rope_theta: float = 100.0, rope_dim_ratio: float = 0.75, bias: bool = True, eps: float = 1e-5, num_register_tokens: int = 4, *, device=None):
        super().__init__()
        dim = heads * dim_head
        self.patch_size, self.patch_size_t, self.out_channels = patch_size, patch_size_t, out_channels
        self.num_register_tokens = num_register_tokens
        self.pos_embed = RotaryEmbeddingND(int(dim_head * rope_dim_ratio), rope_theta, device=device)
        self.x_embedder = nn.Linear(in_channels, dim, device=device)
        self.register_tokens = nn.Parameter(torch.empty(1, num_register_tokens, dim, device=device))
        self.register_buffer("mask_token", torch.empty(1, 1, dim, device=device))
        self.transformer_blocks = nn.ModuleList([TransformerBlock(heads, dim_head, bias, eps, device=device) for _ in range(num_layers)])
        self.norm_out = nn.LayerNorm(dim, eps=eps, device=device)
        self.proj_out = nn.Linear(dim, out_channels * patch_size_t * patch_size * patch_size, device=device)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        batch, _, latent_t, latent_h, latent_w = x.shape
        h = self.x_embedder(x.flatten(2).transpose(1, 2))
        patches = h.shape[1]
        h = torch.cat((h, self.register_tokens.to(h).expand(batch, -1, -1), torch.zeros_like(h[:, :1])), dim=1)
        ids = create_token_ids((latent_t, latent_h, latent_w), x.device, x.dtype).expand(batch, -1, -1)
        ids = torch.cat((ids, torch.zeros(batch, 1 + self.num_register_tokens, 3, device=x.device, dtype=x.dtype)), dim=1)
        rope = self.pos_embed(ids)
        for block in self.transformer_blocks:
            h = block(h, rope)
        output = self.proj_out(self.norm_out(h))[:, :patches]
        output = output.view(batch, latent_t, latent_h, latent_w, self.out_channels, self.patch_size_t, self.patch_size, self.patch_size)
        return output.permute(0, 4, 1, 5, 2, 6, 3, 7).contiguous().reshape(batch, self.out_channels, latent_t * self.patch_size_t, latent_h * self.patch_size, latent_w * self.patch_size)


class MiniMaxH3VideoVAE(nn.Module):
    """Decoder-only H3 VAE. The public ``decode`` contract matches upstream_vae.py."""
    def __init__(self, *, device=None, tiling: bool = True):
        super().__init__()
        self.vae_ratio, self.vae_ratio_t = 16, 4
        self.clip_length, self.token_drop = 17, 3
        self.frame_pre_padding = (-self.clip_length) % self.vae_ratio_t
        self.tokens_chunk_size = math.ceil(self.clip_length / self.vae_ratio_t)
        self.token_overlap = (-self.token_drop) % self.tokens_chunk_size
        self.frame_overlap = max(self.token_overlap * self.vae_ratio_t - self.frame_pre_padding, 0)
        self.tiling, self.tile_size, self.tile_overlap_min = tiling, 256, 64
        self.post_quant_conv = nn.Conv3d(24, 24, 1, device=device)
        self.decoder = ViT3DDecoder(device=device)
        self.register_buffer("latents_mean", torch.tensor(LATENTS_MEAN, device=device))
        self.register_buffer("latents_std", torch.tensor(LATENTS_STD, device=device))
        self.register_buffer("pixel_mean", torch.tensor(IMAGENET_MEAN, device=device).view(1, 3, 1, 1, 1), persistent=False)
        self.register_buffer("pixel_std", torch.tensor(IMAGENET_STD, device=device).view(1, 3, 1, 1, 1), persistent=False)

    @classmethod
    def from_safetensors(cls, path: str | Path, *, device: str | torch.device = "cuda", tiling: bool = True, dtype: torch.dtype = torch.float32) -> "MiniMaxH3VideoVAE":
        model = cls(device="meta", tiling=tiling)
        expected = model.state_dict()
        if os.getenv("H3_FAST_SAFETENSORS", "").lower() in {"1", "true", "yes", "on"}:
            from fastsafetensors import fastsafe_open

            fast_device = "cuda:0" if str(device) == "cuda" else str(device)
            with fastsafe_open(filenames=[str(path)], nogds=True, device=fast_device) as checkpoint:
                available_weights = {
                    name: checkpoint.get_tensor(name).clone().detach()
                    for name in checkpoint.keys()
                }
            available = set(available_weights)
            missing = sorted(set(expected) - available)
            shape_errors = [(name, tuple(expected[name].shape), tuple(available_weights[name].shape)) for name in expected if name in available and tuple(expected[name].shape) != tuple(available_weights[name].shape)]
            if missing or shape_errors:
                details = ([f"missing: {', '.join(missing)}"] if missing else []) + ([f"shape mismatch: {shape_errors}"] if shape_errors else [])
                raise ValueError("incompatible H3 VAE checkpoint; " + "; ".join(details))
            weights = {name: available_weights[name].to(device=device, dtype=dtype) for name in expected}
        elif os.getenv("H3_DISABLE_MMAP", "").lower() in {"1", "true", "yes", "on"}:
            from safetensors.torch import load

            with open(path, "rb") as file:
                available_weights = load(file.read())
            available = set(available_weights)
            missing = sorted(set(expected) - available)
            shape_errors = [(name, tuple(expected[name].shape), tuple(available_weights[name].shape)) for name in expected if name in available and tuple(expected[name].shape) != tuple(available_weights[name].shape)]
            if missing or shape_errors:
                details = ([f"missing: {', '.join(missing)}"] if missing else []) + ([f"shape mismatch: {shape_errors}"] if shape_errors else [])
                raise ValueError("incompatible H3 VAE checkpoint; " + "; ".join(details))
            weights = {name: available_weights[name].to(device=device, dtype=dtype) for name in expected}
        else:
            with safe_open(str(path), framework="pt", device=str(device)) as checkpoint:
                available = set(checkpoint.keys())
                missing = sorted(set(expected) - available)
                shape_errors = [(name, tuple(expected[name].shape), tuple(checkpoint.get_slice(name).get_shape())) for name in expected if name in available and tuple(expected[name].shape) != tuple(checkpoint.get_slice(name).get_shape())]
                if missing or shape_errors:
                    details = ([f"missing: {', '.join(missing)}"] if missing else []) + ([f"shape mismatch: {shape_errors}"] if shape_errors else [])
                    raise ValueError("incompatible H3 VAE checkpoint; " + "; ".join(details))
                weights = {name: checkpoint.get_tensor(name).to(dtype=dtype) for name in expected}
        model.load_state_dict(weights, strict=True, assign=True)
        return model

    def _decode_pixels(self, z: torch.Tensor) -> torch.Tensor:
        return self.decoder(self.post_quant_conv(z))

    def split_tiles(self, length: int) -> tuple[list[int], list[int], list[int]]:
        if self.tile_size >= length:
            return [0], [length], []
        count = math.ceil(length / self.tile_size)
        while self.tile_size * count - self.tile_overlap_min * (count - 1) < length:
            count += 1
        overlaps = [self.tile_overlap_min] * (count - 1)
        for index in range((self.tile_size * count - sum(overlaps) - length) // self.vae_ratio):
            overlaps[index % len(overlaps)] += self.vae_ratio
        starts = [0]
        for overlap in overlaps:
            starts.append(starts[-1] + self.tile_size - overlap)
        return starts, [self.tile_size] * count, overlaps

    @staticmethod
    def blend(a: torch.Tensor, b: torch.Tensor, extent: int, dim: int) -> torch.Tensor:
        extent = min(a.shape[dim], b.shape[dim], extent)

        positions = torch.arange(extent, device=b.device, dtype=b.dtype)
        weight_a = 1 - positions / extent
        weight_b = positions / extent

        shape = [1] * a.ndim
        shape[dim] = extent
        weight_a = weight_a.view(shape)
        weight_b = weight_b.view(shape)

        slice_a = [slice(None)] * a.ndim
        slice_a[dim] = slice(-extent, None)
        slice_b = [slice(None)] * b.ndim
        slice_b[dim] = slice(0, extent)

        blended = a[tuple(slice_a)] * weight_a + b[tuple(slice_b)] * weight_b

        if extent < b.shape[dim]:
            slice_b_rest = [slice(None)] * b.ndim
            slice_b_rest[dim] = slice(extent, None)
            return torch.cat((blended, b[tuple(slice_b_rest)]), dim=dim)
        return blended

    def tiled_decode(self, z: torch.Tensor) -> torch.Tensor:
        height, width = z.shape[-2] * self.vae_ratio, z.shape[-1] * self.vae_ratio
        y_starts, y_lengths, y_overlaps = self.split_tiles(height)
        x_starts, x_lengths, x_overlaps = self.split_tiles(width)
        canvas = None
        row_tails = []
        output_y = 0
        for row_index, (y, tile_height) in enumerate(zip(y_starts, y_lengths)):
            next_row_tails = []
            left_tail = None
            output_x = 0
            for column_index, (x, tile_width) in enumerate(zip(x_starts, x_lengths)):
                tile = self._decode_pixels(z[..., y // self.vae_ratio:(y + tile_height) // self.vae_ratio, x // self.vae_ratio:(x + tile_width) // self.vae_ratio])
                if row_index < len(y_starts) - 1:
                    next_row_tails.append(tile[..., -y_overlaps[row_index]:, :].clone())
                next_left_tail = tile[..., :, -x_overlaps[column_index]:].clone() if column_index < len(x_starts) - 1 else None
                if row_index:
                    tile = self.blend(row_tails[column_index], tile, y_overlaps[row_index - 1], -2)
                if column_index:
                    tile = self.blend(left_tail, tile, x_overlaps[column_index - 1], -1)
                left_tail = next_left_tail
                if row_index < len(y_starts) - 1:
                    tile = tile[..., :-y_overlaps[row_index], :]
                if column_index < len(x_starts) - 1:
                    tile = tile[..., :, :-x_overlaps[column_index]]
                if canvas is None:
                    canvas = torch.empty(*tile.shape[:-2], height, width, dtype=tile.dtype, device=tile.device)
                canvas[..., output_y:output_y + tile.shape[-2], output_x:output_x + tile.shape[-1]].copy_(tile)
                output_x += tile.shape[-1]
            row_tails = next_row_tails
            output_y += tile.shape[-2]
        return canvas

    def _decode_temporal_pad_frames(self, z_len: int, pad_tokens: int) -> int:
        if pad_tokens <= 0:
            return 0
        intra_tail = self.clip_length % self.vae_ratio_t
        if intra_tail == 0:
            return pad_tokens * self.vae_ratio_t
        z_len_before_pad = z_len - pad_tokens
        return sum(intra_tail if (z_len_before_pad + index) % self.tokens_chunk_size == 0 else self.vae_ratio_t for index in range(pad_tokens))

    def _decode_temporal_frame_plan(self, z_len: int, chunks: int, pad_tokens: int) -> int:
        chunk_dec = self.tokens_chunk_size * self.vae_ratio_t
        split_count = int(self.token_drop > 0) + 1
        total_frames, final_overlap_frames = 0, 0
        for index in range(chunks):
            token_start = index * self.tokens_chunk_size
            token_end = token_start + self.tokens_chunk_size + self.token_overlap
            clip_token_len = max(0, min(token_end, z_len) - min(token_start, z_len))
            clip_frame_len = clip_token_len * self.vae_ratio_t
            for split in range(split_count):
                frame_start = split * chunk_dec
                frame_end = min(frame_start + chunk_dec, clip_frame_len)
                part = max(0, frame_end - frame_start - self.frame_pre_padding)
                if split == 0:
                    total_frames += part
                else:
                    final_overlap_frames = part
        return total_frames + final_overlap_frames - self._decode_temporal_pad_frames(z_len, pad_tokens)

    def decode_temporal(self, z: torch.Tensor) -> torch.Tensor:
        chunk_dec = self.tokens_chunk_size * self.vae_ratio_t
        split_count = int(self.token_drop > 0) + 1
        pseudo_tokens = z.shape[2] + self.token_drop
        pad_tokens = (-pseudo_tokens) % self.tokens_chunk_size
        pseudo_tokens += pad_tokens
        chunks = pseudo_tokens // self.tokens_chunk_size - 1
        if chunks < 1:
            pad_tokens += self.tokens_chunk_size
            chunks += 1
        if pad_tokens:
            z = torch.cat((z, z[:, :, -1:].repeat(1, 1, pad_tokens, 1, 1)), dim=2)
        output_frames = self._decode_temporal_frame_plan(z.shape[2], chunks, pad_tokens)
        output = None
        overlap = None
        write_pos = 0

        def write(part: torch.Tensor) -> None:
            nonlocal output, write_pos
            if part.shape[2] <= 0:
                return
            if output is None:
                shape = list(part.shape)
                shape[2] = output_frames
                output = torch.empty(shape, dtype=part.dtype, device=part.device)
            copy_frames = min(part.shape[2], max(0, output.shape[2] - write_pos))
            if copy_frames > 0:
                output[:, :, write_pos:write_pos + copy_frames].copy_(part[:, :, :copy_frames])
                write_pos += copy_frames

        for index in range(chunks):
            clip = self._adaptive_decode(z[:, :, index * self.tokens_chunk_size:index * self.tokens_chunk_size + self.tokens_chunk_size + self.token_overlap])
            for split in range(split_count):
                frame_start = split * chunk_dec
                frame_end = min(frame_start + chunk_dec, clip.shape[2])
                part = clip[:, :, frame_start:frame_end][:, :, self.frame_pre_padding:]
                if split == 0:
                    if overlap is not None:
                        part = self.blend(overlap, part, self.frame_overlap, -3)
                        overlap = None
                    write(part)
                else:
                    overlap = part.contiguous()
            if index == chunks - 1 and overlap is not None:
                write(overlap)
                overlap = None
        if overlap is not None:
            write(overlap)
        if output is None:
            raise RuntimeError("VAE temporal decode produced no frames")
        return output

    def _adaptive_decode(self, z: torch.Tensor) -> torch.Tensor:
        return self.tiled_decode(z) if self.tiling else self._decode_pixels(z)

    def decode(self, z: torch.Tensor) -> torch.Tensor:
        z = z * self.latents_std.view(1, -1, 1, 1, 1).to(z) + self.latents_mean.view(1, -1, 1, 1, 1).to(z)
        decoded = self._adaptive_decode(z) if z.shape[2] == 1 else self.decode_temporal(z)
        if z.shape[2] == 1:
            decoded = decoded[:, :, -1:]
        pixel_std = torch.tensor(IMAGENET_STD, device=decoded.device, dtype=decoded.dtype).view(1, 3, 1, 1, 1)
        pixel_mean = torch.tensor(IMAGENET_MEAN, device=decoded.device, dtype=decoded.dtype).view(1, 3, 1, 1, 1)
        return (decoded.float() * pixel_std.float() + pixel_mean.float()).clamp_(0, 1).mul_(2).sub_(1)