# Created by BVH & VG, Apr - Aug 2025.

import collections
import os

import lovely_tensors
import torch
from custom.vae.action_codec import (
    actions_to_absolute,
    actions_to_relative,
    normalize_actions,
    prepare_action_normalizer,
    unnormalize_actions,
)

from imaginaire.utils import log

lovely_tensors.monkey_patch()

# Ensure these constants are always correct / up2date:
SIGMA_DATA = 1.0
T5_EMB_MAX_LENGTH = 512

# NOTE(bvh): hard cap on cached_lat_language to keep VRAM bounded on
# datasets with effectively unbounded prompt vocabulary (e.g. DROID has a unique
# action description per sample). Each entry is ~1 MB (512 * 1024 bf16 + small
# mask), so 1024 caps it at ~1 GB per rank.
LAT_LANGUAGE_CACHE_MAX_ENTRIES = 1024


class VAE:
    def __init__(self, device, config, pipe):
        '''
        :param config (dict/Any4DConfig/Predict2Video2WorldModelConfig): also contains:
            - model_manager_config (dict/Predict2ModelManagerConfig),
            - pipe_config (dict/Video2WorldPipelineConfig),
            - wandb (dict).
        :param downsample (tuple): (T, H, W).
        '''
        self.device = device
        self.config = config
        self.pipe = pipe

        # Device and precision
        self.tensor_kwargs = dict(device=device, dtype=torch.bfloat16)
        # self.vae_ratio = vae_ratio

        # Video tokenizer (created by Video2WorldPipeline)
        self.video_tokenizer = pipe.tokenizer
        
        # (T, H, W) downsampling factors between pixel & latent space
        self.vae_ratio = (self.video_tokenizer.temporal_compression_factor,
                          self.video_tokenizer.spatial_compression_factor,
                          self.video_tokenizer.spatial_compression_factor)

        # Text tokenizer (created by Video2WorldPipeline)
        self.text_encoder = pipe.text_encoder

        # Dict to store cached language tokens.
        # OrderedDict so we can evict the oldest entry past the size cap (LRU-on-write).
        self.cached_lat_language: "collections.OrderedDict[str, dict]" = collections.OrderedDict()

        # Action parameterization:
        # TODO: revisit this
        self.relative_actions = False
        self.action_normalizer = None
        
        if hasattr(config, 'relative_actions'):
            self.relative_actions = config.relative_actions
            
        if hasattr(config, 'normalizer_stats') and len(config.normalizer_stats) > 0:
            self.action_normalizer = prepare_action_normalizer(
                config.normalizer_stats, config.relative_actions)

        if self.text_encoder is None:
            from custom.utils.io import read_pickle
            _dt_path = getattr(config, 'default_text_path', 'custom/vae/default_text.pkl')
            self.default_text = read_pickle(_dt_path)
            self.default_text = {key: val.to(**self.tensor_kwargs) 
                for key, val in self.default_text.items()}

            log.warning(f'Using custom/vae/default_text.pkl: {self.default_text}...')

        else:
            self.default_text = None

    def _cache_lat_language(self, prompt_key: str, value: dict):
        '''Insert into cached_lat_language with LRU-on-write eviction.
        '''
        cache = self.cached_lat_language
        if prompt_key in cache:
            cache.move_to_end(prompt_key)
            return  # already cached, nothing more to do
        cache[prompt_key] = value
        while len(cache) > LAT_LANGUAGE_CACHE_MAX_ENTRIES:
            cache.popitem(last=False)  # drop the oldest

    @torch.no_grad()
    def encode_text(self, prompt):
        '''
        :param prompt (list of str): Batch of language instructions.
        :return retval (dict).
            t5_text_embeddings: (B, L, C) = (B, 512, 1024) tensor of bfloat16.
            t5_text_mask: (B, L) = (B, 512) tensor of bfloat16.
                Contains some True (1.0) followed by all False (0.0).
        '''
        # Do nothing if there is no text encoder 
        if self.text_encoder is None:
            results = {key: val for key, val in self.default_text.items()}
            results['t5_text_embeddings'] = results['t5_text_embeddings'].repeat(len(prompt), 1, 1)
            results['t5_text_mask'] = results['t5_text_mask'].repeat(len(prompt), 1)
            return results

        # Encode text
        with torch.no_grad():
            log.debug(f'Encoding prompt: {prompt}')  # , rank0_only=False)
            (text_embedding, text_mask) = self.text_encoder.encode_prompts(
                prompt, max_length=T5_EMB_MAX_LENGTH, return_mask=True)
            # ^ (B, L, C) tensor of float32, (B, L) tensor of bool.

        # Return text embeddings and mask
        result = {
            't5_text_embeddings': text_embedding.to(**self.tensor_kwargs),
            't5_text_mask': text_mask.to(**self.tensor_kwargs),
        }
        return result

    @torch.no_grad()
    def encode_text_as_needed(self, sample):
        '''
        X
        '''
        # TODO @vitor: Clean up this method / improve readability
        # TODO @vitor: seems like t5_text_mask can now also be 'pruned',
        # update this method / verify correctness

        data_key = 'anydata' if 'anydata' in sample else 'vidar'
        
        if 'lat_language' in sample[data_key]:
            lat_language = sample[data_key]['lat_language'][(0,0)]
            
            # Only prompt is provided, so encode everything
            if len(lat_language) == 1 or lat_language['t5_text_embeddings'] == 'pruned':
                sample_text = self.encode_text(sample['prompt'])
            
            # Everything is provided
            else:
                
                # Move text data to device and dtype
                sample_text = {key: lat_language[key].to(**self.tensor_kwargs)
                               for key in lat_language
                               if key not in ['prompt']}
                
                # Check if each batch sample is valid (invalid samples were padded with zeros during collation)
                for b in range(len(lat_language['prompt'])):
                    valid = all([val[b].sum() != 0
                                 for key, val in lat_language.items()
                                 if key not in ['prompt']])
                    
                    # For each invalid sample
                    if not valid:
                        # Process that batch sample again to generate valid latents
                        prompt = lat_language['prompt'][b]

                        if prompt in self.cached_lat_language:
                            # Use cached version to save time and memory
                            sample_text_b = self.cached_lat_language[prompt]
                            self.cached_lat_language.move_to_end(prompt)  # LRU touch
                        else:
                            # Process and cache.
                            sample_text_b = self.encode_text(prompt)
                            self._cache_lat_language(prompt, {
                                key: v.detach().clone().contiguous() for key, v in sample_text_b.items()
                            })

                        # Assign processed or cached latents to batch sample
                        for key, val in sample_text_b.items():
                            sample_text[key][b] = val
        
        # If language latents are not provided, encode everything and save caches
        else:
            is_cached = True
            lat_language = []
            prompt = sample['prompt']
            B = len(prompt)
            
            for b in range(B): # Check each prompt
                if prompt[b] in self.cached_lat_language:
                    lat_language.append(self.cached_lat_language[prompt[b]])
                else:
                    is_cached = False
                    break
            
            if is_cached: # If everything is cached
                sample_text = {key: torch.stack([lat_language[i][key] for i in range(B)], 0) for key in lat_language[0].keys()}
            
            else: # Encode and store
                sample_text = self.encode_text(sample['prompt'])
                for b in range(B):
                    # clone() the per-sample slice so it doesn't keep the full
                    # (B, 512, 1024) batched tensor alive via the view's
                    # storage. Was the source of a slow VRAM leak on DROID
                    # (every sample carries a unique action prompt).
                    self._cache_lat_language(prompt[b], {
                        key: val[b].detach().clone().contiguous() for key, val in sample_text.items()
                    })
        
        sample = {**sample, **sample_text}

        # TODO: is this still relevant?
        if 'negative_prompt' in sample:
            sample_negative_text = self.encode_text(sample['negative_prompt'])
            sample_negative_text = {
                'neg_t5_text_embeddings': sample_negative_text['t5_text_embeddings'],
                'neg_t5_text_mask': sample_negative_text['t5_text_mask'],
            }
            sample = {**sample, **sample_negative_text}

        return sample

    @torch.no_grad()
    def encode_single_video(self, raw_video):
        '''
        :param raw_video: (B, 3, Tp, Hp, Wp) tensor of float in [-1, 1].
        :return latent_video: (B, Cl, Tl, Hl, Wl) tensor of float.
        '''
        (B, Cp, Tp, Hp, Wp) = raw_video.shape
        assert Cp == 3, \
            f'Raw video must have 3 channels, but got {Cp}'
        assert (Tp - 1) % self.video_tokenizer.temporal_compression_factor == 0, \
            f'raw Tp - 1 must be divisible by temporal_compression_factor: ' \
            f'{self.video_tokenizer.temporal_compression_factor}, but got {Tp}'

        raw_video = raw_video.to(**self.tensor_kwargs, non_blocking=True)

        # Encode pixel video into latent video.
        # NOTE(bvh): We explicitly multiply by sigma_data here to match the diffusion pipeline.
        # NOTE(bvh): Unlike old Cosmos, no padding and unpadding seems to be necessary anymore.
        # NOTE(bvh): ~245 ms/call rgb, ~490 ms/call plucker (2 sub-encodes) at T=41, H=W=288, A100; see 23.md.
        # HIGH PRI TODO(bvh): parallelize/batch VAE calls to make cams faster
        with torch.no_grad():
            latent_video = self.video_tokenizer.encode(raw_video)
            latent_video = latent_video * SIGMA_DATA

        return latent_video

    @torch.no_grad()
    def encode_single_video_rgb(self, raw_video):
        '''
        :param raw_video: (B, 3, Tp, Hp, Wp) tensor of float in [0, 1].
        :return latent_video: (B, Cl, Tl, Hl, Wl) tensor of float.
        '''
        assert raw_video.shape[1] == 3, 'RGB video must have 3 channels'

        if raw_video.isnan().any():
            nan_per_batch = raw_video.isnan().flatten(1).any(dim=1)
            log.error(f'NaN in RAW RGB before VAE encode! '
                      f'shape={raw_video.shape}, '
                      f'affected batch elements: {nan_per_batch.nonzero().flatten().tolist()}', rank0_only=False)
            raw_video = raw_video.nan_to_num(0.0)

        raw_video = raw_video * 2.0 - 1.0  # Now values are in [-1, 1]
        output = self.encode_single_video(raw_video)
        return output

    @torch.no_grad()
    def encode_single_video_depth(self, raw_video, key=None):
        '''
        :param raw_video: (B, 1, Tp, Hp, Wp) tensor of float in [0, 1].
        :return latent_video: (B, Cl, Tl, Hl, Wl) tensor of float.
        '''
        assert raw_video.shape[1] == 1, 'Depth video must have 1 channel'
        
        # invalid = (raw_video == 0.0)
        raw_video = raw_video * 2.0 - 1.0  # Now values are in [-1, 1]
        # raw_video[invalid] = 0.0

        # if key is not None:
        #     self.invalid_mask[key] = invalid

        raw_video = torch.cat([raw_video] * 3, dim=1)
        output = self.encode_single_video(raw_video)
        return output

    @torch.no_grad()
    def encode_single_video_points(self, raw_video):
        '''
        :param raw_video: (B, 3, Tp, Hp, Wp) tensor of float in [-1, 1].
        :return latent_video: (B, Cl, Tl, Hl, Wl) tensor of float.
        '''
        assert raw_video.shape[1] == 3, 'Points video must have 3 channels'

        output = self.encode_single_video(raw_video)
        return output

    @torch.no_grad()
    def encode_single_video_cams(self, raw_video, with_orig=False):
        '''
        :param raw_video: (B, 6/9, Tp, Hp, Wp) tensor of float in [-1, 1].
        :return latent_video: (B, Cl * 2/3, Tl, Hl, Wl) tensor of float.
        '''
        assert raw_video.shape[1] in [6, 9], 'Plucker video must have 6 (rays, cross) or 9 (rays, cross, origin) channels'

        # NOTE(bvh): Origins are optional; see vidar/geometry/cameras/pinhole.py for more details.
        raw_video_rays  = raw_video[:, 0:3]
        raw_video_cross = raw_video[:, 3:6]
        if with_orig:
            raw_video_orig = raw_video[:, 6:9]

        latent_video_rays = self.encode_single_video(raw_video_rays)
        latent_video_cross = self.encode_single_video(raw_video_cross)
        if with_orig:
            latent_video_orig = self.encode_single_video(raw_video_orig)
        
        if with_orig:
            output = torch.cat([latent_video_rays, latent_video_cross, latent_video_orig], dim=1)
        else:
            output = torch.cat([latent_video_rays, latent_video_cross], dim=1)
        
        return output

    @torch.no_grad()
    def encode_single_video_auto(self, key, raw_video):
        '''
        :param key (str): rgb* / depth* / points* / cams* / ...
        :param raw_video: (B, Cp, Tp, Hp, Wp) tensor of float in [0, 1] or [-1, 1]
            (value range depends on modality).
        :return latent_video: (B, Cl, Tl, Hl, Wl) tensor of float.
        '''
        (B, Cp, Tp, Hp, Wp) = raw_video.shape
        raw_video = raw_video.clone().detach()

        if key.startswith('rgb'):
            latent_video = self.encode_single_video_rgb(raw_video)
        
        elif key.startswith('depth'):
            if Cp == 1:
                # New way (we are presumably using actual values)
                latent_video = self.encode_single_video_depth(raw_video, key=key)

            elif Cp == 3:
                # Old way (we are presumably using colormaps)
                log.warning('Probably using obsolete way to encode depth video as colormaps')
                latent_video = self.encode_single_video_rgb(raw_video)

            else:
                raise ValueError(f'Unknown or invalid number of channels for raw depth video: {Cp}')
            
        elif key.startswith('points'):
            latent_video = self.encode_single_video_points(raw_video)

        elif key.startswith('cams'):
            latent_video = self.encode_single_video_cams(raw_video)

        else:
            raise ValueError(f'Unknown or invalid entry: {key}')

        return latent_video

    @torch.no_grad()
    def encode_single_lowdim_auto(self, key, raw_data):
        '''
        :param key (str): action / proprio / ...
        :param raw_data: (B, T, C) tensor of float.
        :return latent_data: (B, T, C) tensor of float.
        '''
        latent_data = raw_data.clone().detach()
        
        if key.startswith('action') or key.startswith('proprio'):
            if self.relative_actions:
                (latent_data, bases) = actions_to_relative(latent_data)
                
                # TODO @vitor: store this info in batch instead
                # self.extra['absolute_actions'] = bases
            
            if self.action_normalizer is not None:
                latent_data = normalize_actions(latent_data, self.action_normalizer)
            
            latent_data *= self.config.action_multiplier

        elif key.startswith('bbox'):
            pass  # Passthrough for now

        elif key.startswith('traj'):
            pass  # Passthrough for now

        else:
            # log.warning(f'Low-dim entry {key} is not explicitly handled by VAE, '
            #             f'simply passing through from raw to latent as fallback.')
            raise ValueError(f'Unknown or invalid entry: {key}')

        return latent_data

    @torch.no_grad()
    def encode_a4d_batch(self, data_batch, phase: str = 'train'):
        '''
        :param data_batch (dict): Any4D data batch containing entries after data augmentations,
            collate, and model transforms, but with possibly partial or all raw data.
        :param phase (str): train / val / test.
        :return batch_enc (dict): Any4D data batch after VAE encoding.
        '''
        
        # Create shallow copies to avoid modifying the original
        batch_enc = dict(data_batch)
        a4d_raw = dict(data_batch['a4d_raw'])
        a4d_latent = dict(data_batch['a4d_latent'])

        # Process not-yet-encoded language
        batch_enc = self.encode_text_as_needed(batch_enc)

        # Process all other not-yet-encoded entries (high-dim videos & low-dim info)
        for (k, e) in a4d_raw.items():
            assert k not in ['fps', 'prompt', 'num_frames'], \
                f'Key {k} should go in the parent sample dictionary instead'
            
            if k in ['video_dims', 'used_views']:  # Skip metadata important for logistics.
                continue

            elif k in self.config.all_highdim_entries:
                # Check if already pre-encoded, e.g. in case of cached latents
                if k in a4d_latent:
                    continue
                
                latent_video = self.encode_single_video_auto(k, e)
                a4d_latent[k] = latent_video

            elif k in self.config.all_lowdim_entries:
                # Check if already pre-encoded, e.g. in case of cached latents
                if k in a4d_latent:
                    continue

                latent_data = self.encode_single_lowdim_auto(k, e)
                a4d_latent[k] = latent_data

            elif self.config.ignore_unused_data:
                log.debug(f'Ignoring provided but unused entry: {k}: {e.shape}')
            
            else:
                raise ValueError(f'Unknown or invalid entry: {k}')

        batch_enc['a4d_latent'] = a4d_latent

        return batch_enc

    @torch.no_grad()
    def decode_single_video(self, latent_video):
        '''
        :param latent_video: (B, Cl, Tl, Hl, Wl) tensor of float.
        :return raw_video: (B, 3, Tp, Hp, Wp) tensor of float in [-1, 1].
        '''
        # Decode latent video into pixel video.
        # NOTE(bvh): We explicitly divide by sigma_data here to match the diffusion pipeline.
        # NOTE(bvh): Unlike old Cosmos, no padding and unpadding seems to be necessary anymore.
        with torch.no_grad():
            latent_video = latent_video / SIGMA_DATA
            raw_video = self.video_tokenizer.decode(latent_video)

        return raw_video

    @torch.no_grad()
    def decode_single_video_rgb(self, latent_video):
        '''
        :param latent_video: (B, Cl, Tl, Hl, Wl) tensor of float.
        :return raw_video: (B, 3, Tp, Hp, Wp) tensor of float in [0, 1].
        '''
        output = self.decode_single_video(latent_video)  # Returned values are in [-1, 1]
        output = (output + 1.0) / 2.0  # Now values are in [0, 1]
        return output

    @torch.no_grad()
    def decode_single_video_depth(self, latent_video, key=None):
        '''
        :param latent_video: (B, Cl, Tl, Hl, Wl) tensor of float.
        :return raw_video: (B, 1, Tp, Hp, Wp) tensor of float in [0, 1].
        '''
        output = self.decode_single_video(latent_video)  # Returned values are in [-1, 1]
        output = output.mean(dim=1, keepdim=True)
        output = (output + 1.0) / 2.0  # Now values are in [0, 1]
        
        # if key is not None and key in self.invalid_mask.keys():
        #     output[self.invalid_mask[key]] = 0.0
        
        return output

    @torch.no_grad()
    def decode_single_video_points(self, latent_video, key=None):
        '''
        :param latent_video: (B, Cl, Tl, Hl, Wl) tensor of float.
        :return raw_video: (B, 3, Tp, Hp, Wp) tensor of float in [-1, 1].
        '''
        output = self.decode_single_video(latent_video)  # Returned values are in [-1, 1]
        return output

    @torch.no_grad()
    def decode_single_video_cams(self, latent_video, with_orig=False):
        '''
        :param latent_video: (B, 2/3 * Cl, Tl, Hl, Wl) tensor of float.
        :return raw_video: (B, 6/9, Tp, Hp, Wp) tensor of float in [-1, 1].
        '''
        latent_video_rays  = latent_video[:, 0:16]
        latent_video_cross = latent_video[:, 16:32]
        if with_orig:
            latent_video_orig = latent_video[:, 32:48]
        
        output_rays = self.decode_single_video(latent_video_rays)  # Returned values are in [-1, 1]
        output_cross = self.decode_single_video(latent_video_cross)  # Returned values are in [-1, 1]
        if with_orig:
            output_orig = self.decode_single_video(latent_video_orig)  # Returned values are in [-1, 1]

        if with_orig:
            output = torch.cat([output_rays, output_cross, output_orig], dim=1)
        else:
            output = torch.cat([output_rays, output_cross], dim=1)
        
        return output

    @torch.no_grad()
    def decode_single_video_auto(self, key, latent_video, with_orig=False):
        '''
        :param key (str): rgb_* / depth_* / points_* / cams_* / ...
        :param latent_video: (B, Cl, Tl, Hl, Wl) tensor of float.
        :return pixel_video: (B, 1/3/6/9, Tp, Hp, Wp) tensor of float in [0, 1] or [-1, 1]
            (channel count and value range both depend on modality).
        NOTE(bvh): For consistent visualization & metrics, follow this by visuals.py:viz_video_auto().
        '''
        latent_video = latent_video.clone().detach()

        if key.startswith('rgb'):
            pixel_video = self.decode_single_video_rgb(latent_video)
            pixel_video = pixel_video.clamp(0.0, 1.0)
            # ^ (B, 3, Tp, Hp, Wp) tensor of float in [0, 1] with RGB pixels.
            
        elif key.startswith('depth'):
            # TODO(bvh): Support old Kubric depth viz behavior
            pixel_video = self.decode_single_video_depth(latent_video, key)
            pixel_video = pixel_video.clamp(0.001, 1.0) # Small non-zero minimum to avoid invalid predicted pixels
            # ^ (B, 1, Tp, Hp, Wp) tensor of float in [0, 1] with real but scaled depth values.

        elif key.startswith('points'):
            pixel_video = self.decode_single_video_points(latent_video)
            pixel_video = pixel_video.clamp(-1.0, 1.0)
            # ^ (B, 3, Tp, Hp, Wp) tensor of float in [-1, 1] with real but scaled XYZ coordinates.

        elif key.startswith('cams'):
            pixel_video = self.decode_single_video_cams(latent_video, with_orig=with_orig)
            pixel_video = pixel_video.clamp(-1.0, 1.0)
            # ^ (B, 6/9, Tp, Hp, Wp) tensor of float in [-1, 1] with real but scaled Plucker vectors;
            # channel dimension depends on with_orig.
            
        else:
            raise ValueError(f'Unknown or invalid sample key: {key}')

        return pixel_video

    @torch.no_grad()
    def decode_single_lowdim_auto(self, key, latent_data):
        '''
        :param key (str): action / proprio / ...
        :param latent_data: (B, T, C) tensor of float.
        :return raw_data: (B, T, C) tensor of float.
        '''
        raw_data = latent_data.clone().detach()
        
        if key.startswith('action') or key.startswith('proprio'):
            raw_data /= self.config.action_multiplier

            # TODO(bvh) @vitor: this part below is taken from old visuals code, but
            # I believe this order is wrong? does not match encode_single_lowdim_auto()
            
            if self.relative_actions:
                raw_data = actions_to_absolute(raw_data, None)
            
            if self.action_normalizer is not None:
                raw_data = unnormalize_actions(raw_data, self.action_normalizer)
        
        elif key.startswith('bbox'):
            pass  # Passthrough for now

        elif key.startswith('traj'):
            pass  # Passthrough for now

        else:
            # log.warning(f'Low-dim entry {key} is not explicitly handled by VAE, '
            #             f'simply passing through from latent to raw as fallback.')
            raise ValueError(f'Unknown or invalid entry: {key}')
            
        return raw_data

    @torch.no_grad()
    def decode_a4d_batch(self, entries_latent, phase: str = 'train'):
        '''
        :param entries_latent (dict): Any4D entries after DiT forward pass.
        :param phase (str): train / val / test.
        :return entries_raw (dict): Any4D entries after VAE decoding.
        '''
        entries_raw = dict()

        # Process all other not-yet-encoded entries (high-dim videos & low-dim info)
        for (k, e) in entries_latent.items():
            if k in ['video_dims', 'used_views']:  # Skip metadata important for logistics.
                continue

            elif k.endswith('_mask'):
                continue
            
            elif k in self.config.all_highdim_entries:
                latent_video = self.decode_single_video_auto(k, e)
                entries_raw[k] = latent_video

            elif k in self.config.all_lowdim_entries:
                latent_data = self.decode_single_lowdim_auto(k, e)
                entries_raw[k] = latent_data

            elif self.config.ignore_unused_data:
                log.debug(f'Ignoring provided but unused entry: {k}: {e.shape}')

            else:
                raise ValueError(f'Unknown or invalid entry: {k}')

        return entries_raw