from __future__ import annotations

import os
os.environ["TOKENIZERS_PARALLELISM"] = "false"

from pathlib import Path
from tqdm.auto import tqdm

import torch
import torch.nn.functional as F
from torch import cat, nn, Tensor, is_tensor, tensor
from torch.nn import Module
from torch.utils.data import DataLoader
from torch.optim import AdamW
from accelerate import Accelerator

from einops import rearrange, repeat
import einx
from x_transformers import Encoder

from diffusers.models.transformers.transformer_cosmos import CosmosTransformer3DModel, CosmosTransformerBlock
from diffusers.models.autoencoders.autoencoder_kl_cosmos import AutoencoderKLCosmos, CosmosEncoder3d
from transformers import T5EncoderModel, T5TokenizerFast, T5Config

from einops.layers.torch import Reduce
from torch_einops_utils import shape_with_replace, lens_to_mask, masked_mean

# helpers

def exists(v):
    return v is not None

def default(v, d):
    return v if exists(v) else d

def cast_tensor(val, device = None):
    return tensor(val, device = device) if not is_tensor(val) else val

def logit_normal_sample(size, mu = 0.0, sigma = 1.0, device = 'cpu'):
    z = torch.randn(size, device = device) * sigma + mu
    return torch.sigmoid(z)

# action chunk summarizer

class ActionChunkSummarizer(Module):
    def __init__(
        self,
        action_dim,
        dim_out,
        dim = 256,
        depth = 2,
        heads = 4
    ):
        super().__init__()
        self.proj_in = nn.Linear(action_dim, dim)
        self.encoder = Encoder(
            dim = dim,
            depth = depth,
            heads = heads
        )
        self.pool = Reduce('b t d -> b d', 'mean')
        self.proj_out = nn.Linear(dim, dim_out)
        
        nn.init.zeros_(self.proj_out.weight)
        nn.init.zeros_(self.proj_out.bias)

    def forward(self, actions):
        # actions: (b, t, d_act)
        x = self.proj_in(actions)
        x = self.encoder(x)
        # pool sequence: (b, t, d) -> (b, d)
        x = self.pool(x)
        return self.proj_out(x)

# constants

TINY_TRANSFORMER_CONFIG = dict(
    in_channels = 16,
    out_channels = 16,
    num_attention_heads = 1,
    attention_head_dim = 16,
    mlp_ratio = 1.0,
    text_embed_dim = 32,
    adaln_lora_dim = 32,
    patch_size = (1, 2, 2),
    max_size = (16, 64, 64),
    extra_pos_embed_type = None,
    concat_padding_mask = False,
)

TINY_VAE_CONFIG = dict(
    in_channels = 3,
    out_channels = 3,
    latent_channels = 16,
    encoder_block_out_channels = (8, 16),
    decode_block_out_channels = (8, 16),
    temporal_compression_ratio = 1,
    spatial_compression_ratio = 4,
    num_layers = 1,
    attention_resolutions = (),
    resolution = 64,
)

TINY_T5_CONFIG = dict(
    vocab_size = 32128,
    d_model = 32,
    d_kv = 8,
    d_ff = 64,
    num_layers = 1,
    num_heads = 1,
)

REAL_TRANSFORMER_CONFIG = dict(
    in_channels = 16,
    out_channels = 16,
    num_attention_heads = 32,
    attention_head_dim = 128,
    mlp_ratio = 4.0,
    text_embed_dim = 1024,
    patch_size = (1, 2, 2),
    max_size = (128, 240, 240),
    extra_pos_embed_type = "learnable",
    concat_padding_mask = False,
)

REAL_VAE_CONFIG = dict(
    in_channels = 3,
    out_channels = 3,
    latent_channels = 16,
    encoder_block_out_channels = (128, 256, 512, 512),
    decode_block_out_channels = (256, 512, 512, 512),
    temporal_compression_ratio = 8,
    spatial_compression_ratio = 8,
)

REAL_T5_CONFIG = dict(
    vocab_size = 32128,
    d_model = 1024,
    d_kv = 64,
    d_ff = 2048,
    num_layers = 12,
    num_heads = 16,
)

DEFAULT_LORA_CONFIG = dict(
    r = 8,
    lora_alpha = 16,
    target_modules = ["to_q", "to_k", "to_v", "to_out.0"],
    lora_dropout = 0.05,
    bias = "none",
)



class ActionConditionedCosmosTransformer3DModel(CosmosTransformer3DModel):
    def __init__(self, *args, action_dim: int | None = None, **kwargs):
        super().__init__(*args, **kwargs)
        self.action_dim = action_dim
        
        if not exists(action_dim):
            return
            
        hidden_size = self.config.num_attention_heads * self.config.attention_head_dim
        
        self.action_summarizer = ActionChunkSummarizer(
            action_dim = action_dim,
            dim_out = hidden_size
        )

    def forward(
        self,
        hidden_states: torch.Tensor,
        actions: torch.Tensor | None = None,
        **kwargs
    ):
        if not (exists(actions) and exists(self.action_dim)):
            return super().forward(hidden_states, **kwargs)

        orig_patch_embed = self.patch_embed
        
        class WrappedPatchEmbed(Module):
            def __init__(self, patch_embed, summarizer):
                super().__init__()
                self.patch_embed = patch_embed
                self.summarizer = summarizer
                
            def forward(self, x):
                x = self.patch_embed(x)
                
                # Ensure actions match the precision of the summarizer
                actions_cast = actions.to(self.summarizer.proj_in.weight.dtype)
                act_emb = self.summarizer(actions_cast)
                
                if x.ndim == 5:
                    from einops import rearrange
                    act_emb = rearrange(act_emb, 'b d -> b 1 1 1 d')
                elif x.ndim == 3:
                    from einops import rearrange
                    act_emb = rearrange(act_emb, 'b d -> b 1 d')
                
                return x + act_emb.to(x.dtype)

        self.patch_embed = WrappedPatchEmbed(orig_patch_embed, self.action_summarizer)
        
        try:
            out = super().forward(hidden_states, **kwargs)
        finally:
            self.patch_embed = orig_patch_embed
            
        return out

# main class

class CosmosPredictWrapper(Module):
    def __init__(
        self,
        model_name: str = 'nvidia/Cosmos-1.0-Diffusion-7B-Video2World',
        extract_layers: int | list[int] | None = None,
        random_weights: bool = False,
        tiny: bool = False,
        normalize = lambda t: (t - 0.5) * 2.0,
        extract_layer: int | None = None,
        lora_path: str | None = None,
        action_dim: int | None = None,
        torch_dtype: Optional[torch.dtype] = None,
        **kwargs
    ):
        super().__init__()
        extract_layers = default(default(extract_layers, extract_layer), 19)
        self.extract_layers = [extract_layers] if isinstance(extract_layers, int) else extract_layers
        self.return_list = isinstance(extract_layers, list)

        self.hook_handles = []
        self.cached_hidden_states = []
        
        if random_weights:
            self._init_random_weights(tiny = tiny, action_dim = action_dim, torch_dtype = torch_dtype)
        else:
            self._init_pretrained(model_name, action_dim = action_dim, torch_dtype = torch_dtype, **kwargs)

        self.dim_latent = self.transformer.config.num_attention_heads * self.transformer.config.attention_head_dim
        self.normalize = normalize

        self.action_dim = action_dim

        if exists(lora_path):
            self.load_lora(lora_path)

        # latents mean and std initialization
        dtype = default(torch_dtype, torch.bfloat16)

        if  random_weights:
            self.latents_mean = torch.zeros(1, self.vae_latent_channels, 1, 1, 1).to(self.device, dtype=dtype)
            self.latents_std = torch.ones(1, self.vae_latent_channels, 1, 1, 1).to(self.device, dtype=dtype)
        else:
            mean_list = self.vae.config.latents_mean[:self.vae_latent_channels]
            std_list = self.vae.config.latents_std[:self.vae_latent_channels]

            self.latents_mean = rearrange(torch.tensor(mean_list), 'c -> 1 c 1 1 1').to(self.device, dtype=dtype)
            self.latents_std = rearrange(torch.tensor(std_list), 'c -> 1 c 1 1 1').to(self.device, dtype=dtype)

        self._register_hook()

    @property
    def device(self):
        return next(self.parameters()).device

    def _init_pretrained(self, model_name: str, action_dim: int | None = None, torch_dtype: Optional[torch.dtype] = None, **kwargs):
        from diffusers import CosmosVideoToWorldPipeline
        pipeline = CosmosVideoToWorldPipeline.from_pretrained(model_name, **kwargs)
        
        # Load base models
        self.vae, self.transformer, self.text_encoder, self.tokenizer = pipeline.vae, pipeline.transformer, pipeline.text_encoder, pipeline.tokenizer
        
        # Replace with action-conditioned variants if action_dim is provided
        if exists(action_dim):
            self.transformer = ActionConditionedCosmosTransformer3DModel(
                action_dim = action_dim,
                **self.transformer.config
            ).to(self.transformer.device, self.transformer.dtype)

        self.vae_temporal_compression_ratio = self.vae.config.temporal_compression_ratio
        self.vae_spatial_compression_ratio = self.vae.config.spatial_compression_ratio
        self.vae_latent_channels = self.vae.config.latent_channels

        if exists(torch_dtype):
            self.to(torch_dtype)

        del pipeline

    def _init_random_weights(self, tiny: bool = False, action_dim: int | None = None, torch_dtype: Optional[torch.dtype] = None):
        config_t = TINY_TRANSFORMER_CONFIG if tiny else REAL_TRANSFORMER_CONFIG
        config_v = TINY_VAE_CONFIG if tiny else REAL_VAE_CONFIG
        config_5 = TINY_T5_CONFIG if tiny else REAL_T5_CONFIG

        num_layers = max(28 if not tiny else 2, *[layer + 1 for layer in self.extract_layers])
        
        if exists(action_dim):
            self.transformer = ActionConditionedCosmosTransformer3DModel(num_layers = num_layers, action_dim = action_dim, **config_t)
        else:
            self.transformer = CosmosTransformer3DModel(num_layers = num_layers, **config_t)
            
        self.vae = AutoencoderKLCosmos(**config_v)

        self.text_encoder = T5EncoderModel(T5Config(**config_5))
        self.tokenizer = T5TokenizerFast.from_pretrained("google-t5/t5-small")
        
        self.vae_temporal_compression_ratio = config_v['temporal_compression_ratio']
        self.vae_spatial_compression_ratio = config_v['spatial_compression_ratio']
        self.vae_latent_channels = config_v['latent_channels']

        if exists(torch_dtype) or True: # always ensure a default dtype for random weights
            self.to(default(torch_dtype, torch.bfloat16))

    def _get_dense_time(
        self,
        t: torch.Tensor,
        batch: int,
        time_len: int,
        dim: int = 5,
        dtype: torch.dtype = torch.float32
    ):
        """Helper to broadcast time t [B] to [B, 1, T, 1, 1] depending on dim."""
        from einops import repeat
        
        if dim >= 3:
            t_expanded = repeat(t, 'b -> b 1 t 1 1', t=time_len)
        else:
            t_expanded = repeat(t, 'b -> b 1 1 1 1')
            
        return t_expanded.to(dtype)

    def __del__(self):
        for handle in getattr(self, 'hook_handles', []): handle.remove()

    def _register_hook(self):
        for layer_index in self.extract_layers:
            target = self.transformer.transformer_blocks[layer_index]
            self.hook_handles.append(target.register_forward_hook(lambda m, i, o: self.cached_hidden_states.append(o.detach().cpu())))

    def load_lora(self, lora_path: str):
        from peft import PeftModel
        if isinstance(self.transformer, PeftModel):
            self.transformer.load_adapter(lora_path, adapter_name = "default")
        else:
            self.transformer = PeftModel.from_pretrained(self.transformer, lora_path)

    @torch.no_grad()
    def sample_flow_trajectory(
        self,
        latents: Tensor,            
        encoder_states: Tensor,      
        target_tau: float = 1.0,     
        steps: int = 10,            
        num_prefix_frames: int = 0,
        actions: Tensor | None = None
    ) -> Tensor:
        
        batch_size, _, total_frames, _, _ = latents.shape
      
        def get_dense_time(t_value):
            ts = torch.zeros(
                (batch_size, 1, total_frames, 1, 1), 
                device=self.device, 
                dtype=latents.dtype
            )
            ts[:, :, num_prefix_frames:] = t_value
            return ts
        
        if target_tau >= 1.0 - 1e-4:
            ts = get_dense_time(999)
            self.cached_hidden_states.clear()
            kwargs = dict(actions=actions) if exists(actions) and exists(self.action_dim) else dict()
            self.transformer(
                hidden_states=latents,
                encoder_hidden_states=encoder_states,
                timestep=ts,
                return_dict=False,
                **kwargs
            )
            return latents

        # ODE solver
        timesteps = torch.linspace(1.0, target_tau, steps + 1, device=self.device)
        curr_latents = latents.clone()

        pbar = tqdm(range(steps), desc = "Sampling frames", leave = False)
        for i in pbar:
            t_curr = timesteps[i]
            t_next = timesteps[i+1]
            dt = t_next - t_curr  

            ts = get_dense_time(t_curr)

            self.cached_hidden_states.clear() 
            
            kwargs = dict(actions=actions) if exists(actions) and exists(self.action_dim) else dict()
            velocity = self.transformer(
                hidden_states=curr_latents,
                encoder_hidden_states=encoder_states,
                timestep=ts * 1000, 
                return_dict=False,
                **kwargs
            )[0]

            # update future latents only
            vel_future = velocity[:, :, num_prefix_frames:]
            curr_latents[:, :, num_prefix_frames:] += vel_future * dt

        #  Final pass to get hidden states at target_tau
        self.cached_hidden_states.clear()
        
        final_ts = get_dense_time(target_tau)
    
        kwargs = dict(actions=actions) if exists(actions) and exists(self.action_dim) else dict()
        self.transformer(
            hidden_states=curr_latents,
            encoder_hidden_states=encoder_states,
            timestep=final_ts * 1000,
            return_dict=False,
            **kwargs
        )

        return curr_latents

    def forward(
        self,
        videos: Tensor,
        prompts: str | list[str] | None = None,
        prompt_token_ids: Tensor | None = None,
        actions: Tensor | None = None,
        predict_num_future_latents = 0, # number of future frames to predict at inference - presumably the given video will be fixed at T = 0, with the future predicted frames at T = 1
        video_flow_target_tau: float = 0.0,
        inference_steps: int = 15,
        return_decoded_latents: bool = False,
        next_videos: Tensor | None = None,
        return_flow_loss: bool = False
    ) -> Tensor | list[Tensor]:

        batch, *_ = videos.shape
        videos = self.normalize(videos)

        if isinstance(prompts, str): prompts = [prompts] * batch

        self.cached_hidden_states.clear()

        videos = rearrange(videos, 'b t c h w -> b c t h w').to(self.device)

        if exists(prompt_token_ids):
            text_inputs = dict(input_ids = prompt_token_ids.to(self.device))
        else:
            text_inputs = self.tokenizer(default(prompts, [""] * batch), return_tensors = "pt", padding = True, truncation = True, max_length = 512).to(self.device)

        encoder_states = self.text_encoder(**text_inputs)
        if hasattr(encoder_states, 'last_hidden_state'):
            encoder_states = encoder_states.last_hidden_state
        encoder_states = encoder_states.to(self.transformer.dtype)

        # Repeat videos temporally to map 1 VAE latent to 1 frame 
        # (offsetting the VAE's temporal downsampling factor)
        videos = repeat(videos, 'b c t h w -> b c (t r) h w', r = self.vae_temporal_compression_ratio)

        # Actions are simply condition embeddings at the transformer level now
        encode_out = self.vae.encode(videos).latent_dist
        latents = encode_out.sample() if hasattr(encode_out, 'sample') else encode_out
            
        latents = (latents - self.latents_mean.to(latents.device)) / self.latents_std.to(latents.device)

        if return_flow_loss:
            assert exists(next_videos), "next_videos must be provided for flow matching loss"
            
            # Encode Target
            with torch.no_grad():
                next_videos_norm = self.normalize(next_videos)
                next_videos_norm = rearrange(next_videos_norm, 'b t c h w -> b c t h w').to(self.device)
                next_videos_repeated = repeat(next_videos_norm, 'b c t h w -> b c (t r) h w', r = self.vae_temporal_compression_ratio)
                
                encode_out_target = self.vae.encode(next_videos_repeated).latent_dist
                latents_target = encode_out_target.sample() if hasattr(encode_out_target, 'sample') else encode_out_target
                latents_target = (latents_target - self.latents_mean.to(self.device)) / self.latents_std.to(self.device)
            
            num_prefix_latents = latents.shape[2]
            
            # Flow matching noise
            noise = torch.randn_like(latents_target)
            t = torch.rand(batch, device=self.device, dtype=self.transformer.dtype)
            t_expanded = self._get_dense_time(t, batch, latents_target.shape[2], dim=latents_target.ndim, dtype=self.transformer.dtype)
            
            latents_target_noisy = t_expanded * noise + (1.0 - t_expanded) * latents_target
            target_velocity = noise - latents_target
            
            noisy_full_latents = torch.cat((latents, latents_target_noisy), dim=2)
            
            transformer_kwargs = dict()
            if exists(actions) and exists(self.action_dim):
                transformer_kwargs['actions'] = actions
            
            ts_input = self._get_dense_time(t, batch, noisy_full_latents.shape[2], dim=5, dtype=self.transformer.dtype) * 1000.0
            
            pred_velocity = self.transformer(
                hidden_states=noisy_full_latents,
                encoder_hidden_states=encoder_states,
                timestep=ts_input,
                return_dict=False,
                **transformer_kwargs
            )[0]
            
            import torch.nn.functional as F
            pred_velocity_future = pred_velocity[:, :, num_prefix_latents:]
            loss = F.mse_loss(pred_velocity_future, target_velocity)
            return loss

        ending_latents = None

        if predict_num_future_latents > 0:
            # conditioning on time=0 for prefix (clean), and time=999 for future (noise)
            
            num_prefix_frames = latents.shape[2]

            pred_shape = shape_with_replace(latents, {2: predict_num_future_latents})
            future_noise = torch.randn(pred_shape, device = latents.device, dtype = latents.dtype)

            starting_latents = cat((latents, future_noise), dim = 2)

            if exists(actions):
                # We interpolate actions below to match latents, so we don't need a strict assert on length
                # However, it should at least have some length > 0
                assert actions.shape[1] > 0

            ending_latents = self.sample_flow_trajectory(
                latents = starting_latents,
                encoder_states = encoder_states,
                target_tau = video_flow_target_tau,
                steps = inference_steps,
                num_prefix_frames = num_prefix_frames,
                actions = actions
            )
        else:
            if exists(actions):
                assert actions.shape[1] > 0
            
            # extract representations
            ts = torch.full((batch, 1, latents.shape[2], 1, 1), video_flow_target_tau, device = self.device, dtype = latents.dtype)
            
            kwargs = dict(actions=actions) if exists(actions) else dict()
            self.transformer(
                hidden_states=latents,
                encoder_hidden_states=encoder_states,
                timestep=ts * 1000, 
                return_dict=False,
                **kwargs
            )

        if return_decoded_latents and exists(ending_latents):
            # Un-normalize latents
            ending_latents = (ending_latents * self.latents_std.to(ending_latents.device)) + self.latents_mean.to(ending_latents.device)
            # Cosmos VAE decode
            decoded_video = self.vae.decode(ending_latents).sample
            # decode output is typically [-1, 1], lets map to [0, 1]
            decoded_video = (decoded_video / 2 + 0.5).clamp(0, 1)
            # We must subsample the decoded video to reverse the temporal 8x duplication
            decoded_video = decoded_video[:, :, ::self.vae_temporal_compression_ratio]
            # Back from (B, C, T, H, W) to (B, T, C, H, W)
            decoded_video = rearrange(decoded_video, 'b c t h w -> b t c h w')
            return decoded_video

        hiddens = self.cached_hidden_states[:len(self.extract_layers)]

        return hiddens if self.return_list else hiddens[0]

class Cosmos2_5PredictWrapper(CosmosPredictWrapper):
    def __init__(
        self,
        model_name: str = 'nvidia/Cosmos-Predict2.5-2B',
        **kwargs
    ):
        super().__init__(
            model_name = model_name,
            **kwargs
        )

# dynamics trainer

class DynamicsTrainer(Module):
    def __init__(
        self,
        model,
        dataset,
        batch_size = 1,
        lr = 1e-4,
        weight_decay = 0.,
        max_grad_norm = 0.5,
        shuffle = True,
        cpu = False,
        accelerate_kwargs: dict = {}
    ):
        super().__init__()
        self.accelerator = Accelerator(cpu = cpu, **accelerate_kwargs)
        self.model = model
        
        self.dataloader = DataLoader(
            dataset, 
            batch_size = batch_size, 
            shuffle = shuffle
        )
        
        self.optimizer = AdamW(model.parameters(), lr = lr, weight_decay = weight_decay)
        
        self.model, self.optimizer, self.dataloader = self.accelerator.prepare(
            self.model,
            self.optimizer,
            self.dataloader
        )
        
        self.max_grad_norm = max_grad_norm

    def train_step(
        self,
        videos,
        next_videos,
        actions = None,
        prompt_token_ids = None
    ):
        self.model.train()
        
        with self.accelerator.autocast():
            loss = self.model(
                videos = videos,
                next_videos = next_videos,
                actions = actions,
                prompt_token_ids = prompt_token_ids,
                return_flow_loss = True
            )

        self.accelerator.backward(loss)
        
        if exists(self.max_grad_norm):
            self.accelerator.clip_grad_norm_(self.model.parameters(), self.max_grad_norm)
        
        self.optimizer.step()
        self.optimizer.zero_grad()
        
        return loss.item()

    def train(
        self,
        num_steps = 1000
    ):
        self.model.train()
        
        loader_iter = iter(self.dataloader)
        
        pbar = tqdm(range(num_steps), desc = 'Fine-tuning Dynamics Model')
        
        for _ in pbar:
            try:
                batch = next(loader_iter)
            except StopIteration:
                loader_iter = iter(self.dataloader)
                batch = next(loader_iter)
            
            videos = batch.get('video')
            next_videos = batch.get('next_video')
            actions = batch.get('actions')
            prompt_token_ids = batch.get('prompt_token_ids')
            
            loss = self.train_step(
                videos = videos,
                next_videos = next_videos,
                actions = actions,
                prompt_token_ids = prompt_token_ids
            )
            
            pbar.set_postfix(loss = f'{loss:.4f}')
            
        self.accelerator.print('Fine-tuning complete.')
