# Wan 2.1 VAE adapter exposing the minimal TokenizerInterface used by the rest of Any4D.
# Wraps cosmos_predict2.tokenizers.tokenizer.VAE (which already inlines the Wan 2.1 VAE
# with the same Wan-published mean/std scaling) so we don't duplicate model code here.

import os

import torch
import torch.nn as nn

from cosmos_predict2.tokenizers.tokenizer import VAE as _CosmosWanVAE
from imaginaire.utils import log


# Canonical Wan 2.1 VAE spec.
WAN_LATENT_CH = 16
WAN_TEMPORAL_COMPRESSION = 4
WAN_SPATIAL_COMPRESSION = 8


def _resolve_vae_path(vae_pth: str) -> str:
    """Resolve `vae_pth` to a local file path.

    Accepts:
      - empty string: returns "" (caller leaves VAE randomly initialized)
      - a local file path: returns as-is
      - a local directory containing one of the canonical file names: returns the file inside
      - an `s3://...` URL: downloads via `get_checkpoint` (rank-0 download + barrier) and
        returns the cached local path. The cosmos VAE loader does not have an S3 backend
        registered, so we materialize the file ourselves first.
    """
    if not vae_pth:
        return vae_pth
    if vae_pth.startswith('s3://'):
        from custom.utils.utils import get_checkpoint
        vae_pth = get_checkpoint(vae_pth)
    if os.path.isfile(vae_pth):
        return vae_pth
    candidates = ['Wan2.1_VAE.pth', 'Wan2.1-VAE.pth', 'vae.pt', 'wan_vae.pth']
    for c in candidates:
        p = os.path.join(vae_pth, c)
        if os.path.isfile(p):
            return p
    raise FileNotFoundError(f'Cannot find Wan VAE weights in {vae_pth} (tried: {candidates})')


def _strip_model_prefix(state_dict):
    """Strip the `model.` prefix that some Wan VAE releases include in their state dict."""
    if 'model_state' in state_dict:
        state_dict = state_dict['model_state']
    return {k[len('model.'):] if k.startswith('model.') else k: v for k, v in state_dict.items()}


class WanVAEAdapter(nn.Module):
    """Minimal TokenizerInterface-like wrapper around the Wan 2.1 VAE.

    Required attributes: latent_ch, temporal_compression_factor, spatial_compression_factor.
    Required methods: encode(video) -> latent, decode(latent) -> video, reset_dtype().
    """

    def __init__(
        self,
        vae_pth: str = "",
        chunk_duration: int = 41,
        name: str = "wan_vae",
        load_mean_std: bool = False,
        dtype: str = "bfloat16",
    ):
        super().__init__()
        self.name = name
        self.chunk_duration = chunk_duration
        self.latent_ch = WAN_LATENT_CH
        self.temporal_compression_factor = WAN_TEMPORAL_COMPRESSION
        self.spatial_compression_factor = WAN_SPATIAL_COMPRESSION
        self.load_mean_std = load_mean_std  # unused — Wan uses fixed Wan-published scaling
        self.vae_pth = vae_pth
        self.latent_chunk_duration = (chunk_duration - 1) // WAN_TEMPORAL_COMPRESSION + 1

        self._dtype = {
            'float32': torch.float32,
            'float16': torch.float16,
            'bfloat16': torch.bfloat16,
        }[dtype]

        resolved_pth = _resolve_vae_path(vae_pth) if vae_pth else ''
        self._vae = _CosmosWanVAE(
            z_dim=WAN_LATENT_CH,
            vae_pth=resolved_pth,
            dtype=self._dtype,
            device='cuda',
            is_amp=False,
            temporal_window=WAN_TEMPORAL_COMPRESSION,
        )
        if not resolved_pth:
            log.warning('[WanVAEAdapter] No vae_pth provided — Wan VAE is randomly initialized')
        self._vae.model.eval()

    def reset_dtype(self):
        self._vae.model = self._vae.model.to(dtype=self._dtype)
        self._vae.dtype = self._dtype

    def to(self, *args, **kwargs):
        """Forward .to(device=..., dtype=...) to the underlying VAE nn.Module.

        Cosmos's `VAE` is a plain Python class (not an nn.Module), so the standard
        `nn.Module.to` inherited from `WanVAEAdapter` would not move its weights.
        Override here so callers can use the adapter as a drop-in tokenizer.
        """
        self._vae.model = self._vae.model.to(*args, **kwargs)
        return self

    @torch.no_grad()
    def encode(self, video: torch.Tensor) -> torch.Tensor:
        """:param video: (B, 3, T, H, W) in [-1, 1]. :return latent: (B, 16, Tl, Hl, Wl)."""
        return self._vae.encode(video.to(dtype=self._dtype)).to(dtype=self._dtype)

    @torch.no_grad()
    def decode(self, latent: torch.Tensor) -> torch.Tensor:
        return self._vae.decode(latent.to(dtype=self._dtype)).to(dtype=self._dtype)