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

Mirrors the encoder half of ``upstream_vae.py`` so keyframe/reference images can
be encoded without ComfyUI. The encoder runs in FP32 and the latent moments are
upcast to FP32 before mean/std normalization (the reference contract).

Causal-conv semantics: spatial padding is reflect; temporal padding is causal
(front-only zeros) with a stride grid that starts at the first input frame
(Comfy autopad "same" / ``causal``). For a single input frame the temporal
taps of the kernel are truncated (Comfy ``autopad="causal_zero"``) so the frame
is not convolved against zero frames.
"""

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 _causal_front_padding(t_in: int, kernel: int, stride: int) -> int:
    """Front-only zero padding matching the reference causal 3D conv.

    ``padding = kernel - 1 - (t_out - 1) * stride`` with ``t_out = ceil(t_in /
    stride)`` (the reference computes the output length from the unpadded
    input). The result is non-negative: with a stride-s grid and a (2k+1)
    kernel, ``ceil(t/stride) >= 1 + (t-1) // stride`` for all t, so
    ``(t_in - 1) % stride * stride >= (kernel - 1) % (2 * stride)``.
    """
    t_out = math.ceil(t_in / stride)
    return max(0, kernel - 1 - (t_out - 1) * stride)


class _CausalConv3d(nn.Module):
    """3D conv: reflect spatial padding, causal (front-zero) temporal padding."""

    def __init__(self, in_channels: int, out_channels: int, kernel_size: int, stride: int | tuple[int, int, int] = 1, spatial_padding: int = 0):
        super().__init__()
        self.kernel_size = kernel_size
        self.stride = stride if isinstance(stride, tuple) else (stride, stride, stride)
        self.spatial_padding = spatial_padding
        self.conv = nn.Conv3d(in_channels, out_channels, kernel_size, stride=self.stride, padding=(0, spatial_padding, spatial_padding))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        t = x.shape[2]
        if t == 1:
            # A single input frame never convolves against zero frames: the
            # temporal taps are truncated (Comfy autopad="causal_zero").
            if self.spatial_padding > 0:
                half = (self.kernel_size - 1) // 2
                kernel = self.conv.weight[:, :, half:half + 2 * self.spatial_padding + 1]
                return F.conv3d(x, kernel, self.conv.bias, (1, self.stride[1], self.stride[2]), (0, self.spatial_padding, self.spatial_padding))
            return self.conv(x)
        front = _causal_front_padding(t, self.kernel_size, self.stride[0])
        if front > 0:
            x = F.pad(x, (0, 0, 0, 0, front, 0))
        return self.conv(x)


class TemporalIsolatedGroupNorm(nn.GroupNorm):
    """GroupNorm with statistics computed per frame (time merged into batch)."""

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if x.dim() != 5:
            return super().forward(x)
        b, c, t, h, w = x.shape
        x = x.permute(0, 2, 1, 3, 4).contiguous().view(b * t, c, 1, h, w)
        x = super().forward(x)
        return x.view(b, t, c, h, w).permute(0, 2, 1, 3, 4).contiguous()


def group_norm_3d(num_channels: int) -> TemporalIsolatedGroupNorm:
    return TemporalIsolatedGroupNorm(num_groups=32, num_channels=num_channels, eps=1e-6, affine=True)


class Downsample3D(nn.Module):
    def __init__(self, in_channels: int, out_channels: int, time_stride: int = 1, space_stride: int = 2):
        super().__init__()
        self.space_stride = space_stride
        self.conv = _CausalConv3d(in_channels, out_channels, kernel_size=3, stride=(time_stride, space_stride, space_stride), spatial_padding=0)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if self.space_stride == 2:
            x = F.pad(x, (0, 1, 0, 1, 0, 0), mode="reflect")
        return self.conv(x)


class ResnetBlock3D(nn.Module):
    def __init__(self, in_channels: int, out_channels: int | None = None):
        super().__init__()
        self.in_channels = in_channels
        self.out_channels = in_channels if out_channels is None else out_channels
        self.norm1 = group_norm_3d(in_channels)
        self.norm2 = group_norm_3d(self.out_channels)
        self.conv1 = _CausalConv3d(in_channels, self.out_channels, kernel_size=3, spatial_padding=1)
        self.conv2 = _CausalConv3d(self.out_channels, self.out_channels, kernel_size=3, spatial_padding=1)
        if self.in_channels != self.out_channels:
            self.nin_shortcut = _CausalConv3d(self.in_channels, self.out_channels, kernel_size=1, spatial_padding=0)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        h = self.conv1(F.silu(self.norm1(x), inplace=True))
        h = self.conv2(F.silu(self.norm2(h), inplace=True))
        if self.in_channels != self.out_channels:
            x = self.nin_shortcut(x)
        return h.add_(x)


class EncoderFCN3D(nn.Module):
    def __init__(self, ch: int, ch_mult: tuple[int, ...], space_down: tuple[int, ...], time_down: tuple[int, ...], num_res_blocks: int, in_channels: int, z_channels: int, double_z: bool = True):
        super().__init__()
        self.num_levels = len(ch_mult)
        self.num_res_blocks = [num_res_blocks] * self.num_levels
        block_mid = [ch * ch_mult[i] for i in range(self.num_levels)]
        block_in = [block_mid[0]] + block_mid[:-1]
        self.conv_in = _CausalConv3d(in_channels, block_in[0], kernel_size=3, spatial_padding=1)
        self.down = nn.ModuleList()
        for i_level in range(self.num_levels):
            down = nn.Module()
            down.block = nn.ModuleList(
                ResnetBlock3D(block_in[i_level] if i == 0 else block_mid[i_level], block_mid[i_level])
                for i in range(self.num_res_blocks[i_level])
            )
            if space_down[i_level] * time_down[i_level] > 1:
                down.downsample = Downsample3D(block_mid[i_level], block_mid[i_level], time_stride=time_down[i_level], space_stride=space_down[i_level])
            self.down.append(down)
        self.norm_out = group_norm_3d(block_mid[-1])
        self.conv_out = _CausalConv3d(block_mid[-1], 2 * z_channels if double_z else z_channels, kernel_size=3, spatial_padding=1)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        h = self.conv_in(x)
        for i_level in range(self.num_levels):
            for i_block in range(self.num_res_blocks[i_level]):
                h = self.down[i_level].block[i_block](h)
            if hasattr(self.down[i_level], "downsample"):
                h = self.down[i_level].downsample(h)
        h = F.silu(self.norm_out(h), inplace=True)
        return self.conv_out(h)


class MiniMaxH3VideoVAEEncoder(nn.Module):
    """Encoder-only H3 VAE. ``encode`` matches the public contract of upstream_vae.py."""

    def __init__(self, *, device: torch.device | str | None = None, tiling: bool = True):
        super().__init__()
        self.vae_ratio, self.vae_ratio_t = 16, 4
        self.clip_length, self.token_drop = 17, 3
        self.tiling, self.tile_size, self.tile_overlap_min = tiling, 256, 64
        self.encoder = EncoderFCN3D(ch=128, ch_mult=(1, 2, 2, 4, 4, 8), space_down=(2, 2, 2, 2, 1, 1), time_down=(1, 2, 2, 1, 1, 1), num_res_blocks=2, in_channels=3, z_channels=24, double_z=True)
        self.quant_conv = nn.Conv3d(48, 48, 1, device=device)
        self.register_buffer("latents_mean", torch.tensor(LATENTS_MEAN), persistent=False)
        self.register_buffer("latents_std", torch.tensor(LATENTS_STD), persistent=False)
        self.register_buffer("pixel_mean", torch.tensor(IMAGENET_MEAN).view(1, 3, 1, 1, 1), persistent=False)
        self.register_buffer("pixel_std", torch.tensor(IMAGENET_STD).view(1, 3, 1, 1, 1), persistent=False)

    @classmethod
    def from_safetensors(cls, path: str | Path, *, device: str | torch.device = "cuda", tiling: bool = True) -> "MiniMaxH3VideoVAEEncoder":
        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)
            available = set()
            with fastsafe_open(filenames=[str(path)], nogds=True, device=fast_device) as checkpoint:
                available = set(checkpoint.keys())
                missing = sorted(set(expected) - available)
                weights = {name: checkpoint.get_tensor(name).clone().detach() for name in expected}
            if missing:
                raise ValueError(f"incompatible H3 VAE checkpoint; missing: {', '.join(missing)}")
            weights = {name: t.to(device=device, dtype=torch.float32) for name, t in weights.items()}
        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)
            if missing:
                raise ValueError(f"incompatible H3 VAE checkpoint; missing: {', '.join(missing)}")
            weights = {name: available_weights[name].to(device=device, dtype=torch.float32) for name in expected}
        else:
            with safe_open(str(path), framework="pt", device=str(device)) as checkpoint:
                available = set(checkpoint.keys())
                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 shape_errors:
                    raise ValueError(f"incompatible H3 VAE checkpoint; shape mismatch: {shape_errors}")
                weights = {name: checkpoint.get_tensor(name).to(dtype=torch.float32) for name in expected}
                missing = sorted(set(expected) - available)
        if missing:
            raise ValueError(f"incompatible H3 VAE checkpoint; missing: {', '.join(missing)}")
        model.load_state_dict(weights, strict=True, assign=True)
        return model

    def _encode_moments(self, x: torch.Tensor) -> torch.Tensor:
        return self.quant_conv(self.encoder(x))

    def _adaptive_encode(self, x: torch.Tensor) -> torch.Tensor:
        if self.tiling:
            return self.tiled_encode(x)
        return self._encode_moments(x)

    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)] * a.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_encode(self, x: torch.Tensor) -> torch.Tensor:
        height, width = x.shape[-2], x.shape[-1]
        y_idx, y_len, y_overlap = self.split_tiles(height)
        x_idx, x_len, x_overlap = self.split_tiles(width)
        rows = [[self._encode_moments(x[..., i_pos:i_pos + i_len, j_pos:j_pos + j_len]) for j_pos, j_len in zip(x_idx, x_len)] for i_pos, i_len in zip(y_idx, y_len)]
        latent_y_overlap = [o // self.vae_ratio for o in y_overlap]
        latent_x_overlap = [o // self.vae_ratio for o in x_overlap]
        result_rows = []
        for i, row in enumerate(rows):
            result_row = []
            for j, tile in enumerate(row):
                if i > 0:
                    tile = self.blend(rows[i - 1][j], tile, latent_y_overlap[i - 1], dim=-2)
                if j > 0:
                    tile = self.blend(row[j - 1], tile, latent_x_overlap[j - 1], dim=-1)
                if i < len(rows) - 1:
                    tile = tile[..., :-latent_y_overlap[i], :]
                if j < len(row) - 1:
                    tile = tile[..., :, :-latent_x_overlap[j]]
                result_row.append(tile)
            result_rows.append(torch.cat(result_row, dim=-1))
        return torch.cat(result_rows, dim=-2)

    def encode_temporal(self, x: torch.Tensor) -> torch.Tensor:
        if x.shape[2] % self.clip_length != 0:
            pad_size = (-x.shape[2]) % self.clip_length
            x = torch.cat([x, x[:, :, -1:].repeat(1, 1, pad_size, 1, 1)], dim=2)
        num_chunks = x.shape[2] // self.clip_length
        z_list = [self._adaptive_encode(x[:, :, i * self.clip_length:(i + 1) * self.clip_length, :, :]) for i in range(num_chunks)]
        z = torch.cat(z_list, dim=2)
        if self.token_drop > 0:
            z = z[:, :, :-self.token_drop]
        return z

    def encode(self, x: torch.Tensor) -> torch.Tensor:
        """``[B, 3, T, H, W]`` pixels in ``[-1, 1]`` -> normalized latents ``[B, 24, T_lat, H//16, W//16]``."""
        if x.ndim == 4:
            x = x.unsqueeze(2)
        x = x.add(1.0).mul_(0.5).sub_(self.pixel_mean.to(x)).div_(self.pixel_std.to(x))
        if x.shape[2] == 1:
            moments = self._adaptive_encode(x)[:, :, -1:, :, :]
        else:
            moments = self.encode_temporal(x)
        mean = torch.chunk(moments.float(), 2, dim=1)[0]
        latents_mean = self.latents_mean.view(1, -1, 1, 1, 1).to(mean)
        latents_std = self.latents_std.view(1, -1, 1, 1, 1).to(mean)
        return (mean - latents_mean) / latents_std