# Created by BVH, May 2026.
# Shape-preserving fake VAE for data-only debug mode (Any4DConfig.data_only_disable_model).
# Encodes via avg_pool3d + channel tile, decodes via slice + nearest upsample.
# Mirrors the real custom.vae.a4d_vae.VAE interface so visuals/metrics still run.

import collections

import torch
import torch.nn.functional as F

from custom.utils.constants import VAE_RATIO_THW
from custom.vae.a4d_vae import VAE as RealVAE, SIGMA_DATA


# Match the shape of real VAE latents: (B, 16, T', H', W').
LATENT_C = 16


class FakeVideoTokenizer:
    '''Drop-in for cosmos tokenizer: only encode/decode + compression factors.'''

    def __init__(self):
        self.temporal_compression_factor = VAE_RATIO_THW[0]
        self.spatial_compression_factor = VAE_RATIO_THW[1]
        assert VAE_RATIO_THW[1] == VAE_RATIO_THW[2], \
            f'FakeVideoTokenizer assumes square spatial compression, got {VAE_RATIO_THW}'

    @torch.no_grad()
    def encode(self, raw_video):
        '''
        :param raw_video: (B, 3, Tp, Hp, Wp) tensor of float in [-1, 1].
        :return latent_video: (B, 16, Tl, Hl, Wl) tensor of float.
        '''
        (Tr, Hr, Wr) = VAE_RATIO_THW
        (B, C, Tp, Hp, Wp) = raw_video.shape
        assert C == 3, f'Expected 3 input channels, got {C}'

        # Match real tokenizer: Tp - 1 must be divisible by Tr; latent Tl = (Tp - 1) // Tr + 1.
        # avg_pool3d with kernel/stride = ratios on the trailing Tp frames gives Tl - 1 latent frames,
        # plus an extra "seed" frame derived from frame 0 to keep the (Tl - 1) * Tr + 1 invariant.
        assert (Tp - 1) % Tr == 0, f'Tp - 1 must be divisible by {Tr}, got Tp={Tp}'
        seed = F.avg_pool3d(
            raw_video[:, :, 0:1],
            kernel_size=(1, Hr, Wr), stride=(1, Hr, Wr))
        body = F.avg_pool3d(
            raw_video[:, :, 1:],
            kernel_size=(Tr, Hr, Wr), stride=(Tr, Hr, Wr))
        latent_rgb = torch.cat([seed, body], dim=2)  # (B, 3, Tl, Hl, Wl)

        # Tile 3 -> 16 channels mod-3: [r,g,b,r,g,b,...] = ceil(16/3) = 6 copies, then slice.
        latent = latent_rgb.repeat(1, (LATENT_C + 2) // 3, 1, 1, 1)[:, :LATENT_C]
        return latent

    @torch.no_grad()
    def decode(self, latent_video):
        '''
        :param latent_video: (B, 16, Tl, Hl, Wl) tensor of float.
        :return raw_video: (B, 3, Tp, Hp, Wp) tensor of float in [-1, 1].
        '''
        (Tr, Hr, Wr) = VAE_RATIO_THW
        (B, C, Tl, Hl, Wl) = latent_video.shape
        assert C >= 3, f'Expected >= 3 latent channels, got {C}'

        # Slice first 3 channels then nearest-upsample. Tp = (Tl - 1) * Tr + 1 to match real VAE.
        rgb_latent = latent_video[:, :3]
        Tp = (Tl - 1) * Tr + 1
        Hp = Hl * Hr
        Wp = Wl * Wr
        raw_video = F.interpolate(
            rgb_latent, size=(Tp, Hp, Wp), mode='nearest')
        return raw_video


class VAE(RealVAE):
    '''Lightweight fake VAE: skips diffusion-VAE / text-encoder loading entirely.
    Used when Any4DConfig.data_only_disable_model is True. Mirrors RealVAE methods
    via inherited code paths; only the video_tokenizer is swapped.
    '''

    def __init__(self, device, config, pipe=None):
        '''
        :param pipe: ignored (kept for signature compat with RealVAE).
        '''
        # NOTE(bvh): bypass super().__init__ on purpose -- avoid touching pipe.tokenizer
        # (heavy VAE weights) and pipe.text_encoder (T5 weights). Set only what the
        # inherited helper methods need.
        self.device = device
        self.config = config
        self.pipe = pipe

        self.tensor_kwargs = dict(device=device, dtype=torch.bfloat16)

        self.video_tokenizer = FakeVideoTokenizer()
        self.vae_ratio = (
            self.video_tokenizer.temporal_compression_factor,
            self.video_tokenizer.spatial_compression_factor,
            self.video_tokenizer.spatial_compression_factor,
        )

        # Force the inherited encode_text() to use the default_text.pkl fallback
        # instead of running a real T5 encoder.
        self.text_encoder = None
        from custom.utils.io import read_pickle
        self.default_text = read_pickle('custom/vae/default_text.pkl')
        self.default_text = {key: val.to(**self.tensor_kwargs)
            for key, val in self.default_text.items()}

        self.cached_lat_language = collections.OrderedDict()

        # Action parameterization: disabled in fake mode so we don't need the normalizer
        # stats file. encode_single_lowdim_auto / decode_single_lowdim_auto still scale
        # by config.action_multiplier (cheap), which is enough for visuals.
        self.relative_actions = False
        self.action_normalizer = None