"""Fast H3 VAE (non-causal): drop-in replacement for MiniMaxH3VideoVAE.decode_base on whole clips.

    from fast_h3_vae import FastH3VAEDecoder
    dec = FastH3VAEDecoder(vae, "fast_h3_vae_noncausal_24l.safetensors")      # vae: the released H3 VAE (for its
    pixels = dec.decode(z)            # z: normalized latents (B=1, T, 24, H, W)   post_quant_conv / embedder / stats)
    rgb_u8 = dec.decode_rgb_uint8(z)  # (F, 16H, 16W, 3) uint8 on the GPU

* any clip length T (t = 0 must be the start of a native 5-latent block, as for the released decoder). The output has
  the same frames as decode_base: 17 frames per full block, the first latent of every block keeps 1 of its 4 frames
  (tested T = 3..15: frame count equal to decode_base, fast vs reference 60.7-61.4 dB). For T = 1 and 2 the released
  decode_base raises an error; this decoder returns 1 and 5 frames by the same rule;
* engine="fast" (default): NonCausalFast3 (FlashAttention-4 varlen windows, fused Triton kernels, one CUDA graph per
  (T, H, W)); engine="ref": the plain PyTorch NonCausalViTDecoder (any GPU with SDPA);
* the fast engine is built on first use of a shape and cached (`max_engines` most recent shapes).
"""
from __future__ import annotations

from collections import OrderedDict

import torch
from safetensors.torch import load_file

from noncausal_vit_decoder import KEEP, NonCausalViTDecoder


class FastH3VAEDecoder:
    def __init__(self, vae, weights: str, *, layers: int = 24, ffn_keep: int | None = None, time_norm: float = 5,
                 engine: str = "fast", max_engines: int = 4, device="cuda", offload_vae: bool = True):
        """offload_vae: after building the decoder, move the released VAE (encoder + 36-layer decoder, ~10 GB) to the CPU;
        decoding only needs its processor (pixel de-normalisation), which works on GPU tensors."""
        self.vae = vae
        m = NonCausalViTDecoder(vae, KEEP[layers], time_norm=time_norm, ffn_keep=ffn_keep)
        r = m.load_state_dict(load_file(weights, device="cpu"), strict=False)
        assert not r.missing_keys and not r.unexpected_keys, r
        self.model = m.to(device).eval().requires_grad_(False)
        if offload_vae:
            vae.to("cpu")
            torch.cuda.empty_cache()
        self.engine, self.max_engines = engine, max_engines
        self._engines: OrderedDict = OrderedDict()

    def _get(self, T, H, W):
        key = (T, H, W)
        if key in self._engines:
            self._engines.move_to_end(key)
            return self._engines[key]
        from noncausal_fast3 import NonCausalFast3
        eng = NonCausalFast3(self.model, T, H, W, cuda_graph=True)
        self._engines[key] = eng
        while len(self._engines) > self.max_engines:
            self._engines.popitem(last=False)
            torch.cuda.empty_cache()
        return eng

    @torch.no_grad()
    def decode(self, z):
        """z: normalized latents (1, T, 24, H, W). Returns decode_base-space pixels (1, 3, F, 16H, 16W), fp32."""
        assert z.dim() == 5 and z.shape[0] == 1 and z.shape[2] == 24, z.shape
        if self.engine == "ref":
            with torch.autocast("cuda", dtype=torch.bfloat16):
                return self.model(z.float()).float()
        _, T, _, H, W = z.shape
        return self._get(T, H, W).decode(z.float().contiguous())

    @torch.no_grad()
    def decode_rgb_uint8(self, z):
        """Same as decode, then the VAE's revert_tensor, clamp and uint8: (F, 16H, 16W, 3) on the GPU."""
        px = self.decode(z)
        return self.vae.processor.revert_tensor(px).clamp_(0, 1).mul_(255).round_().to(torch.uint8)[0].permute(1, 2, 3, 0)