# Created by BVH, Apr 2025.

import torch

from imaginaire.lazy_config import LazyCall as L
from cosmos_predict2.tokenizers.tokenizer import TokenizerInterface
from cosmos_predict2.auxiliary.text_encoder import CosmosT5TextEncoder

from custom.vae.action_codec import actions_to_relative, normalize_actions, prepare_action_normalizer

try:
    from externals.vidar.vidar.utils.viz import viz_depth
except ImportError:
    viz_depth = None

from custom.utils.utils import get_checkpoint, get_folder


# Constants (might need to find a better place for them)
sigma_data = 1.0
t5_embedding_max_length = 512

# Tokenizers taken from Video2WorldPipeline (cosmos_predict2/pipelines/video2world.py)
# TODO: remove pipeline tokenizers so they are not duplicated (right now they are not used, but some attributes are still referenced)

class VAE(torch.nn.Module):
    def __init__(self, device, vae_path, text_encoder_path, downsample=(4,8,8), local_root=''):
        super().__init__()

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

        # Video tokenizer
        self.video_tokenizer = TokenizerInterface(
            chunk_duration=25,
            load_mean_std=False,
            name="tokenizer",
            vae_pth=get_checkpoint(vae_path, local_root=local_root),
        )

        # Text tokenizer
        self.text_encoder = CosmosT5TextEncoder(
            device=device,
            cache_dir=get_folder(text_encoder_path, local_root=local_root),
        )

        # Dict to store cached language tokens
        self.cached_lat_language = {}

        # self.relative_actions = config.relative_actions
        # self.action_normalizer = prepare_action_normalizer(
        #     config.normalizer_stats, config.relative_actions)

    @torch.no_grad()
    def encode_text(self, prompt):

        if isinstance(prompt, torch.Tensor):
            encoded_text = prompt[0,0]
        else:
            # Encode text
            with torch.no_grad():
                text_embedding, text_mask = self.text_encoder.encode_prompts(
                    prompt[0], max_length=512, return_mask=True)

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

    @torch.no_grad()
    def encode_single_video(self, raw_video):
        '''
        :param raw_video: (B, Cp, 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
        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'
        
        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):
        '''
        :param raw_video: (B, 6/9, Tp, Hp, Wp) tensor of float in [-1, 1].
        :return latent_video: (B, Cl * 2, 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'

        raw_video_rays  = raw_video[:, 0:3]
        raw_video_cross = raw_video[:, 3:6]
        # NOTE(bvh): We are ignoring origin for now; see vidar/geometry/cameras/pinhole.py for more details.

        latent_video_rays = self.encode_single_video(raw_video_rays)
        latent_video_cross = self.encode_single_video(raw_video_cross)
        
        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].
        :return latent_video: (B, Cl, Tl, Hl, Wl) tensor of float.
        '''
        (B, Cp, Tp, Hp, Wp) = raw_video.shape

        if key.startswith('rgb'):
            if Cp == self.video_tokenizer.latent_ch:
                # Latents are provided directly
                latent_video = raw_video
            else:
                latent_video = self.encode_single_video_rgb(raw_video)
        
        elif key.startswith('depth'):
            if Cp == self.video_tokenizer.latent_ch:
                # Latents are provided directly
                latent_video = raw_video

            elif 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)
                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 sample key: {key}')

        return latent_video

    @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, Cp, 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
        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):
        '''
        :param latent_video: (B, 2 * Cl, Tl, Hl, Wl) tensor of float.
        :return raw_video: (B, 6, Tp, Hp, Wp) tensor of float in [-1, 1].
        '''
        latent_video_rays  = latent_video[:, 0:16]
        latent_video_cross = latent_video[:, 16:32]
        
        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]
        
        output = torch.cat([output_rays, output_cross], dim=1)
        return output

    @torch.no_grad()
    def decode_single_video_auto(self, key, latent_video, for_viz=False):
        '''
        :param key (str): rgb_* / depth_* / points_* / cams_* / ...
        :param latent_video: (B, Cl, Tl, Hl, Wl) tensor of float.
        :param for_viz (bool): Whether to visualize the video.
        :return raw_video: (B, 6, Tp, Hp, Wp) tensor of float in [0, 1] or [-1, 1].
        NOTE(bvh): We ALWAYS use range [-1, 1] if for_viz = True, for consistent visualization & metrics.
        '''
        if key.startswith('rgb'):
            pixel_video = self.decode_single_video_rgb(latent_video)  # Returned values are in [0, 1]
            pixel_video = pixel_video.clamp(0.0, 1.0)
            
            if for_viz:
                pixel_video = pixel_video * 2.0 - 1.0

        elif key.startswith('depth'):
            # TODO(bvh): Support old Kubric depth viz behavior
            pixel_video = self.decode_single_video_depth(latent_video, key)  # Returned values are in [0, 1]
            pixel_video = pixel_video.clamp(0.0, 1.0)
            
            if for_viz:
                pixel_video = [torch.tensor(
                    viz_depth(pixel_video[0, :, i].float()),
                    dtype=pixel_video.dtype, device=pixel_video.device).permute(2, 0, 1) 
                    for i in range(pixel_video.shape[2])]
                pixel_video = torch.stack(pixel_video, dim=1)[None]
                pixel_video = pixel_video * 2.0 - 1.0
        
        elif key.startswith('points'):
            pixel_video = self.decode_single_video_points(latent_video)  # Returned values are in [-1, 1]
            pixel_video = pixel_video.clamp(-1.0, 1.0)

        elif key.startswith('cams'):
            pixel_video = self.decode_single_video_cams(latent_video)  # Returned values are in [-1, 1]
            pixel_video = pixel_video.clamp(-1.0, 1.0)
            
            if for_viz:
                (pixel_video_rays, pixel_video_cross) = (pixel_video[:, 0:3], pixel_video[:, 3:6])

                viz_type = 'rays'
                if viz_type == 'rays':
                    pixel_video = pixel_video_rays
                elif viz_type == 'comb':
                    # TODO(bvh): Make this a cool raster thingy instead of vertical lines.
                    pixel_video_comb = pixel_video_rays.clone()
                    pixel_video_comb[..., 0::4] = pixel_video_cross[..., 0::4]
                    pixel_video_comb[..., 1::4] = pixel_video_cross[..., 1::4]
                    pixel_video = pixel_video_comb
                else:
                    raise ValueError(f'Invalid viz_type {viz_type}')

                viz_type = self.config.cam_viz_type

                if viz_type == 'rays':
                    pixel_video = pixel_video_rays
                
                elif viz_type == 'comb':
                    # NOTE(bvh): I want to visualize rays and cross together here:
                    pixel_video_comb = pixel_video_rays.clone()
                    w = 16
                    for i in range(w // 2):
                        pixel_video_comb[..., i::w] = pixel_video_cross[..., i::w]
                    pixel_video = pixel_video_comb
                
                else:
                    raise ValueError(f'Invalid viz_type {viz_type}')
                
            # NOTE(bvh): pixel_video now has either 3 or 6 channels depending on for_viz.
            
        else:
            raise ValueError(f'Unknown or invalid sample key: {key}')

        return pixel_video


    @torch.no_grad()
    def process_matrix4d(self, batch, phase: str = 'train', **kwargs):

        tensor_kwargs = dict(dtype=torch.bfloat16)

        # Initialize sample with batch information
        sample = {**batch}

        # Process text
        if 'lat_language' in sample['anydata']:
            lat_language = sample['anydata']['lat_language'][(0,0)]
            # Only prompt is provided, so encode everything
            if len(lat_language) == 1:
                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(**tensor_kwargs) for key in lat_language if key not in ['prompt']}
                # Check if each batch sample is valid (invalid samples were paded 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]
                        else:
                            # Process and cache
                            sample_text_b = self.encode_text([prompt])
                            self.cached_lat_language[prompt] = sample_text_b
                        # 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
        else:
            sample_text = self.encode_text([sample['prompt']])
        sample = {**sample, **sample_text}

        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}

        # Process videos
        for (k, v) in sample['a4d'].items():
            if 'prompt' in k:
                continue

            # elif k == 'fps':
            #     sample[k] = v // self.video_tokenizer.temporal_compression_factor

            if k == 'num_frames':
                assert (v - 1) % self.downsample[0] == 0, \
                    f'raw num_frames - 1 must be divisible by temporal_compression_factor: {self.downsample[0]}, but got {v}'
                sample[k] = (v - 1) // self.downsample[0] + 1

            elif k.endswith('mask'):
                # TODO(bvh): Check self-consistency within each spacetime cube
                if len(v.shape) == 5:
                    sample[k] = v[:, :, 0::self.downsample[0],
                                           self.downsample[1] // 2::self.downsample[1],
                                           self.downsample[2] // 2::self.downsample[2]]
                else:
                    sample[k] = v

            elif k.startswith('action'):
                if self.relative_actions:
                    v, bases = actions_to_relative(v)
                    
                    # TODO(vitor): store this info in batch instead?
                    self.extra['absolute_actions'] = bases
                
                if self.action_normalizer is not None:
                    v = normalize_actions(v, self.action_normalizer)
                sample[k] = v

            else:
                # High-dimensional catch-all:
                if 'rgb' in k and 'lat_rgb' in sample['anydata']:
                    latent_video = sample['anydata']['lat_rgb'][(0,0)].to(**tensor_kwargs)
                else:
                    latent_video = self.encode_single_video_auto(k, v)
                sample[k] = latent_video
                (B, _, Tl, Hl, Wl) = latent_video.shape

        # Ensure that Cosmos-related parameters are present.
        # NOTE(bvh): This requires at least one RGB key to exist, such that Hl and Wl are defined.
        if 'image_size' not in sample:
            sample['image_size'] = torch.tensor([[Hl, Wl, Hl, Wl]], **tensor_kwargs)
        if 'num_frames' not in sample:
            sample['num_frames'] = torch.tensor([Tl], **tensor_kwargs)
        if 'padding_mask' not in sample:
            sample['padding_mask'] = torch.zeros((1, 1, Hl, Wl), **tensor_kwargs)
        if 'train' in phase:
            assert 'loss_mask' in sample, \
                'loss_mask must be manually provided, as it is harder to impute'

        # Necessary for vanilla cosmos
        sample['video_latent'] = sample['rgb_dst']

        # Return original batch with additional information
        return sample