"""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 a keyframe is
never convolved against zero frames.

Weights are read straight from the checkpoint's ``encoder.*`` / ``quant_conv.*``
keys (a direct name-for-name copy into plain tensors) and applied by the
stateless kernels below.
"""

from __future__ import annotations

import math
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)

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
VAE_RATIO = 16
NIN_LEVELS = frozenset({1, 3, 5})
DOWNSAMPLE_LEVELS = frozenset({0, 1, 2, 3})


def _causal_conv3d(x, weight, bias, *, kernel_size, stride, spatial_padding, temporal_causal):
    """Causal 3D conv (matches upstream_vae.CausalConv3d).

    - ``spatial_padding > 0``: reflect H and W by ``spatial_padding`` on each side.
    - ``temporal_causal``: front-zero T by ``kernel_size - 1`` (the reference's
      ``causal_padding[0] * 2`` for ``causal_padding[0]=1``, which is every
      3D-causal conv in the H3 VAE).  A single-frame input truncates the
      temporal taps to the center 1 instead of convolving zero rows.
    - If neither applies: no padding at all (reference early-return).
    """
    if x.shape[2] == 1:
        # Keyframe path (matches reference's `autopad="causal_zero"`):
        # apply spatial-reflect pad (if any), then run an effective 2D conv
        # by slicing the kernel to its center temporal tap (T_out stays 1).
        if spatial_padding > 0:
            x = F.pad(x, (spatial_padding, spatial_padding, spatial_padding, spatial_padding, 0, 0), mode="reflect")
        half = (kernel_size - 1) // 2
        kernel_5d = weight[:, :, half:half + 1, :, :]
        return F.conv3d(x, kernel_5d, bias, (1, stride[1], stride[2]), (0, 0, 0))
    if spatial_padding > 0:
        x = F.pad(x, (spatial_padding, spatial_padding, spatial_padding, spatial_padding, 0, 0), mode="reflect")
    if temporal_causal:
        x = F.pad(x, (0, 0, 0, 0, kernel_size - 1, 0))
    return F.conv3d(x, weight, bias, stride, (0, 0, 0))


def _group_norm_3d(x, weight, bias):
    """GroupNorm (32 groups, eps 1e-6) with per-frame statistics."""
    b, c, t, h, w = x.shape
    y = F.group_norm(x.permute(0, 2, 1, 3, 4).contiguous().view(b * t, c, 1, h, w), 32, weight, bias, 1e-6)
    return y.view(b, t, c, h, w).permute(0, 2, 1, 3, 4).contiguous()


def _resnet(x, p):
    # nin_shortcut uses CausalConv3d(k=1, padding=1) in the reference.
    residual = x if p["nin"] is None else _causal_conv3d(x, p["nin"][0], p["nin"][1], kernel_size=1, stride=(1, 1, 1), spatial_padding=1, temporal_causal=True)
    h = _causal_conv3d(F.silu(_group_norm_3d(x, p["norm1_w"], p["norm1_b"])), p["conv1_w"], p["conv1_b"], kernel_size=3, stride=(1, 1, 1), spatial_padding=1, temporal_causal=True)
    h = _causal_conv3d(F.silu(_group_norm_3d(h, p["norm2_w"], p["norm2_b"])), p["conv2_w"], p["conv2_b"], kernel_size=3, stride=(1, 1, 1), spatial_padding=1, temporal_causal=True)
    return h.add_(residual)


def _downsample(x, p):
    if p["space"] == 2:
        # Reference Downsample3D pads H and W by +1 reflect before the conv.
        x = F.pad(x, (1, 1, 1, 1, 0, 0), mode="reflect")
    # Conv uses padding=(1,0,0) -> causal_padding=(1,0,0), so spatial pad=0,
    # temporal front-zero is applied.
    return _causal_conv3d(x, p["w"], p["b"], kernel_size=3, stride=(p["time"], p["space"], p["space"]), spatial_padding=0, temporal_causal=True)


def _encoder_run(x, E):
    h = _causal_conv3d(x, E["conv_in"][0], E["conv_in"][1], kernel_size=3, stride=(1, 1, 1), spatial_padding=1, temporal_causal=True)
    for level in E["down"]:
        for blk in level["blocks"]:
            h = _resnet(h, blk)
        if level["down"] is not None:
            h = _downsample(h, level["down"])
    h = F.silu(_group_norm_3d(h, E["norm_out_w"], E["norm_out_b"]))
    return _causal_conv3d(h, E["conv_out"][0], E["conv_out"][1], kernel_size=3, stride=(1, 1, 1), spatial_padding=1, temporal_causal=True)


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

    Weights are plain tensors loaded from the checkpoint by canonical name into
    ``self.W`` (a dict), so no ``nn.Module`` sub-hierarchy is needed.
    """

    def __init__(self, *, tiling: bool = True):
        super().__init__()
        self.vae_ratio, self.vae_ratio_t = VAE_RATIO, 4
        self.clip_length, self.token_drop = 17, 3
        self.frame_pre_padding = (-self.clip_length) % self.vae_ratio_t
        self.tiling, self.tile_size, self.tile_overlap_min = tiling, 256, 64
        self.quant_conv = nn.Conv3d(48, 48, 1)
        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)

    def _required_encoder_names(self) -> list[str]:
        names = ["encoder.conv_in.weight", "encoder.conv_in.bias", "encoder.norm_out.weight", "encoder.norm_out.bias", "encoder.conv_out.weight", "encoder.conv_out.bias"]
        for i in range(len(CH_MULT)):
            for b in range(NUM_RES_BLOCKS):
                base = f"encoder.down.{i}.block.{b}."
                names += [
                    base + "conv1.weight", base + "conv1.bias",
                    base + "conv2.weight", base + "conv2.bias",
                    base + "norm1.weight", base + "norm1.bias",
                    base + "norm2.weight", base + "norm2.bias",
                ]
                if b == 0 and i in NIN_LEVELS:
                    names += [base + "nin_shortcut.weight", base + "nin_shortcut.bias"]
            if i in DOWNSAMPLE_LEVELS:
                names += [f"encoder.down.{i}.downsample.conv.weight", f"encoder.down.{i}.downsample.conv.bias"]
        return names

    @classmethod
    def from_safetensors(cls, path: str | Path, *, device: str | torch.device = "cuda", tiling: bool = True) -> "MiniMaxH3VideoVAEEncoder":
        model = cls(tiling=tiling)
        names = model._required_encoder_names()
        with safe_open(str(path), framework="pt", device=str(device)) as ck:
            available = set(ck.keys())
            missing = [n for n in names if n not in available]
            if missing:
                raise ValueError(f"incompatible H3 VAE checkpoint; missing: {', '.join(sorted(missing)[:16])}")
            W = {n: ck.get_tensor(n).to(dtype=torch.float32).to(device) for n in names}

        # Build the structured params dict.
        down = []
        for i in range(len(CH_MULT)):
            mid = CH * CH_MULT[i]
            blocks = []
            for b in range(NUM_RES_BLOCKS):
                base = f"encoder.down.{i}.block.{b}."
                blk = {
                    "conv1_w": W[base + "conv1.weight"], "conv1_b": W[base + "conv1.bias"],
                    "conv2_w": W[base + "conv2.weight"], "conv2_b": W[base + "conv2.bias"],
                    "norm1_w": W[base + "norm1.weight"], "norm1_b": W[base + "norm1.bias"],
                    "norm2_w": W[base + "norm2.weight"], "norm2_b": W[base + "norm2.bias"],
                    "nin": None,
                }
                if base + "nin_shortcut.weight" in W:
                    blk["nin"] = (W[base + "nin_shortcut.weight"], W[base + "nin_shortcut.bias"])
                blocks.append(blk)
            level_down = None
            if i in DOWNSAMPLE_LEVELS:
                ds = f"encoder.down.{i}.downsample.conv."
                level_down = {"w": W[ds + "weight"], "b": W[ds + "bias"], "time": TIME_DOWN[i], "space": SPACE_DOWN[i]}
            down.append({"blocks": blocks, "down": level_down})

        E = {
            "conv_in": (W["encoder.conv_in.weight"], W["encoder.conv_in.bias"]),
            "down": down,
            "norm_out_w": W["encoder.norm_out.weight"], "norm_out_b": W["encoder.norm_out.bias"],
            "conv_out": (W["encoder.conv_out.weight"], W["encoder.conv_out.bias"]),
        }
        model.W = E
        model.quant_conv.to(device, torch.float32)
        for b in ("latents_mean", "latents_std", "pixel_mean", "pixel_std"):
            getattr(model, b).to(device)
        return model

    @torch.inference_mode()
    def _encode_moments(self, x: torch.Tensor) -> torch.Tensor:
        return F.conv3d(_encoder_run(x.to(torch.float32), self.W), self.quant_conv.weight, self.quant_conv.bias)

    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,H,W]`` or ``[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 + 1.0) * 0.5
        x = (x - self.pixel_mean.to(x)) / self.pixel_std.to(x)
        if x.shape[2] == 1:
            # Pad with frame_pre_padding zero frames so the 17-tap stride-2
            # temporal downsampling produces exactly one latent frame (1+3-4).
            pad = self.frame_pre_padding
            x = torch.cat([torch.zeros_like(x[:, :, :pad]), x], dim=2)
            moments = self._adaptive_encode(x)
        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