from typing import List, Optional
import torch
from utils.wan_wrapper import WanDiffusionWrapper, WanTextEncoder, WanVAEWrapper
from utils.memory import gpu, get_cuda_free_memory_gb, move_model_to_device_with_memory_preservation
from pipeline.causal_inference import CausalInferencePipeline
import torch.distributed as dist

class SwitchCausalInferencePipeline(CausalInferencePipeline):

    def __init__(self, args, device, *, generator: WanDiffusionWrapper | None=None, text_encoder: WanTextEncoder | None=None, vae: WanVAEWrapper | None=None):
        super().__init__(args, device, generator=generator, text_encoder=text_encoder, vae=vae)
        self.global_sink = getattr(args, 'global_sink', False)

    def _recache_after_switch(self, output, current_start_frame, new_conditional_dict):
        if not self.global_sink:
            for block_idx in range(self.num_transformer_blocks):
                cache = self.kv_cache1[block_idx]
                cache['k'].zero_()
                cache['v'].zero_()
        for blk in self.crossattn_cache:
            blk['k'].zero_()
            blk['v'].zero_()
            blk['is_init'] = False
        inner = self.generator.model
        if getattr(inner, '_ei_prev_window_start', None) is not None:
            inner._ei_prev_window_start = None
        if current_start_frame == 0:
            return
        num_recache_frames = current_start_frame if self.local_attn_size == -1 else min(self.local_attn_size, current_start_frame)
        recache_start_frame = current_start_frame - num_recache_frames
        frames_to_recache = output[:, recache_start_frame:current_start_frame]
        if frames_to_recache.device.type == 'cpu':
            target_device = next(self.generator.parameters()).device
            frames_to_recache = frames_to_recache.to(target_device)
        batch_size = frames_to_recache.shape[0]
        print(f'num_recache_frames: {num_recache_frames}, recache_start_frame: {recache_start_frame}, current_start_frame: {current_start_frame}')
        device = frames_to_recache.device
        block_mask = self.generator.model._prepare_blockwise_causal_attn_mask(device=device, num_frames=num_recache_frames, frame_seqlen=self.frame_seq_length, num_frame_per_block=self.num_frame_per_block, local_attn_size=self.local_attn_size)
        context_timestep = torch.ones([batch_size, recompute_frames], device=device, dtype=torch.int64) * self.args.context_noise
        self.generator.model.block_mask = block_mask
        with torch.no_grad():
            self.generator(noisy_image_or_video=frames_to_recompute, conditional_dict=new_conditional_dict, timestep=context_timestep, kv_cache=self.kv_cache1, crossattn_cache=self.crossattn_cache, current_start=recompute_start_frame * self.frame_seq_length)
        for blk in self.crossattn_cache:
            blk['k'].zero_()
            blk['v'].zero_()
            blk['is_init'] = False

    def inference(self, noise: torch.Tensor, *, text_prompts_first: List[str], text_prompts_second: List[str], switch_frame_index: int, initial_latent: Optional[torch.Tensor]=None, return_latents: bool=False, low_memory: bool=False):
        batch_size, num_output_frames, num_channels, height, width = noise.shape
        assert num_output_frames % self.num_frame_per_block == 0
        num_blocks = num_output_frames // self.num_frame_per_block
        cond_first = self.text_encoder(text_prompts=text_prompts_first)
        cond_second = self.text_encoder(text_prompts=text_prompts_second)
        if low_memory:
            gpu_memory_preservation = get_cuda_free_memory_gb(gpu) + 5
            move_model_to_device_with_memory_preservation(self.text_encoder, target_device=gpu, preserved_memory_gb=gpu_memory_preservation)
        output_device = torch.device('cpu') if low_memory else noise.device
        output = torch.zeros([batch_size, num_output_frames, num_channels, height, width], device=output_device, dtype=noise.dtype)
        local_attn_cfg = getattr(self.args.model_kwargs, 'local_attn_size', -1)
        kv_policy = ''
        if local_attn_cfg != -1:
            kv_cache_size = local_attn_cfg * self.frame_seq_length
            kv_policy = f'int->local, size={local_attn_cfg}'
        else:
            kv_cache_size = num_output_frames * self.frame_seq_length
            kv_policy = 'global (-1)'
        print(f'kv_cache_size: {kv_cache_size} (policy: {kv_policy}, frame_seq_length: {self.frame_seq_length}, num_output_frames: {num_output_frames})')
        self._initialize_kv_cache(batch_size, dtype=noise.dtype, device=noise.device, kv_cache_size_override=kv_cache_size)
        self._initialize_crossattn_cache(batch_size=batch_size, dtype=noise.dtype, device=noise.device)
        num_input_frames = initial_latent.shape[1] if initial_latent is not None else 0
        current_start_frame = 0
        self.generator.model.local_attn_size = self.local_attn_size
        print(f'[inference] local_attn_size set on model: {self.generator.model.local_attn_size}')
        self._set_all_modules_max_attention_size(self.local_attn_size)
        all_num_frames = [self.num_frame_per_block] * num_blocks
        using_second = False
        for current_num_frames in all_num_frames:
            if not using_second and current_start_frame >= switch_frame_index:
                self._recache_after_switch(output, current_start_frame, cond_second)
                cond_in_use = cond_second
                using_second = True
                print('switch_frame_index', switch_frame_index)
                print('current_start_frame', current_start_frame)
            else:
                cond_in_use = cond_second if using_second else cond_first
            noisy_input = noise[:, current_start_frame - num_input_frames:current_start_frame + current_num_frames - num_input_frames]
            for index, current_timestep in enumerate(self.denoising_step_list):
                timestep = torch.ones([batch_size, current_num_frames], device=noise.device, dtype=torch.int64) * current_timestep
                if index < len(self.denoising_step_list) - 1:
                    _, denoised_pred = self.generator(noisy_image_or_video=noisy_input, conditional_dict=cond_in_use, timestep=timestep, kv_cache=self.kv_cache1, crossattn_cache=self.crossattn_cache, current_start=current_start_frame * self.frame_seq_length)
                    next_timestep = self.denoising_step_list[index + 1]
                    noisy_input = self.scheduler.add_noise(denoised_pred.flatten(0, 1), torch.randn_like(denoised_pred.flatten(0, 1)), next_timestep * torch.ones([batch_size * current_num_frames], device=noise.device, dtype=torch.long)).unflatten(0, denoised_pred.shape[:2])
                else:
                    _, denoised_pred = self.generator(noisy_image_or_video=noisy_input, conditional_dict=cond_in_use, timestep=timestep, kv_cache=self.kv_cache1, crossattn_cache=self.crossattn_cache, current_start=current_start_frame * self.frame_seq_length)
            output[:, current_start_frame:current_start_frame + current_num_frames] = denoised_pred.to(output.device)
            context_timestep = torch.ones_like(timestep) * self.args.context_noise
            self.generator(noisy_image_or_video=denoised_pred, conditional_dict=cond_in_use, timestep=context_timestep, kv_cache=self.kv_cache1, crossattn_cache=self.crossattn_cache, current_start=current_start_frame * self.frame_seq_length)
            current_start_frame += current_num_frames
        video = self.vae.decode_to_pixel(output.to(noise.device), use_cache=False)
        video = (video * 0.5 + 0.5).clamp(0, 1)
        if return_latents:
            return (video, output)
        return video