import torch.nn.functional as F
from .samplers import TTMGuider
from .utils import add_noise_to_reference_video


# Copied from ComfyUI Wanvideo Wrapper
class EncodeWanVideo:
    @classmethod
    def INPUT_TYPES(s):
        return {"required": {
                    "vae": ("VAE",),
                    "image": ("IMAGE",),
                    "enable_vae_tiling": ("BOOLEAN", {"default": False, "tooltip": "Drastically reduces memory use but may introduce seams"}),
                    "tile_x": ("INT", {"default": 272, "min": 64, "max": 2048, "step": 1, "tooltip": "Tile size in pixels, smaller values use less VRAM, may introduce more seams"}),
                    "tile_y": ("INT", {"default": 272, "min": 64, "max": 2048, "step": 1, "tooltip": "Tile size in pixels, smaller values use less VRAM, may introduce more seams"}),
                    "tile_stride_x": ("INT", {"default": 144, "min": 32, "max": 2048, "step": 32, "tooltip": "Tile stride in pixels, smaller values use less VRAM, may introduce more seams"}),
                    "tile_stride_y": ("INT", {"default": 128, "min": 32, "max": 2048, "step": 32, "tooltip": "Tile stride in pixels, smaller values use less VRAM, may introduce more seams"}),
                    },
                    "optional": {
                        "noise_aug_strength": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 10.0, "step": 0.1, "tooltip": "Strength of noise augmentation, helpful for leapfusion I2V where some noise can add motion and give sharper results"}),
                        "latent_strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.1, "tooltip": "Additional latent multiplier, helpful for leapfusion I2V where lower values allow for more motion"}),
                        "mask": ("MASK"),
                    }
                }

    RETURN_TYPES = ("LATENT",)
    RETURN_NAMES = ("reference_latents",)
    FUNCTION = "encode"
    CATEGORY = "Wan22 TimeToMove"

    def encode(self, vae, image, enable_vae_tiling, tile_x, tile_y, tile_stride_x, tile_stride_y, noise_aug_strength=0.0, latent_strength=1.0, mask=None):
        image = image.clone()

        if image.shape[-1] == 4:
            image = image[..., :3]

        if noise_aug_strength > 0.0:
            image = add_noise_to_reference_video(image, ratio=noise_aug_strength)
        
        if enable_vae_tiling:
            latents = vae.encode_tiled(image[:,:,:,:3] * 2.0 - 1.0, tile_size=(tile_x//vae.upscale_ratio, tile_y//vae.upscale_ratio), tile_stride=(tile_stride_x//vae.upscale_ratio, tile_stride_y//vae.upscale_ratio))
        else:
            latents = vae.encode(image[:,:,:,:3] * 2.0 - 1.0)

        if latent_strength != 1.0:
            latents *= latent_strength
                
        print(f"WanVideo Encode: Encoded latents shape {latents.shape}")
 
        return ({"samples": latents, "noise_mask": mask},)


# Copied from ComfyUI Wanvideo Wrapper
class TTMLatentAdd:
    @classmethod
    def INPUT_TYPES(s):
        return {"required": {
                    "latent": ("LATENT", {"tooltip": "wanvideo latent"}),
                    "reference_latents": ("LATENT", {"tooltip": "Reference image to encode"}),
                    "ttm_start_step": ("INT", {"default": 0, "min": 0, "max": 1000, "step": 1, "tooltip": "Start step to apply TTM latent guide"}),
                    "ttm_end_step": ("INT", {"default": 3, "min": 1, "max": 1000, "step": 1, "tooltip": "The step to stop applying TTM"}),
                    "ref_masks": ("MASK", {"tooltip": "Reference mask to encode"}),
                }
        }

    RETURN_TYPES = ("LATENT",)
    RETURN_NAMES = ("latent",)
    FUNCTION = "add"
    CATEGORY = "Wan22 TimeToMove"

    def add(self, latent, reference_latents, ttm_start_step, ttm_end_step, ref_masks):  
              
        if ttm_end_step < max(0, ttm_start_step):
            raise ValueError(f"`ttm_end_step` ({ttm_end_step}) must be >= `ttm_start_step` ({ttm_start_step}).")
        
        mask_sampled = ref_masks[::4]
        mask_sampled = mask_sampled.unsqueeze(1).unsqueeze(0)  # [1, T, 1, H, W]
        
        vae_upscale_factor = 8
        if reference_latents["samples"].shape[1] == 48:
            vae_upscale_factor = 16
        
        # Upsample spatially to latent resolution
        H_latent = mask_sampled.shape[-2] // vae_upscale_factor
        W_latent = mask_sampled.shape[-1] // vae_upscale_factor
        mask_latent = F.interpolate(
            mask_sampled.float(),
            size=(mask_sampled.shape[2], H_latent, W_latent),
            mode="nearest"
        )

        latent["ttm_reference_latents"] = reference_latents["samples"]
        latent["ttm_mask"] =  mask_latent.movedim(2, 1)
        latent["ttm_start_step"] = ttm_start_step
        latent["ttm_end_step"] = ttm_end_step
        
        return (latent,)


class TimeToMoveGuider:    
    @classmethod
    def INPUT_TYPES(s):
        return {"required":
                    {"model": ("MODEL", ),
                    "positive": ("CONDITIONING", ),
                    "negative": ("CONDITIONING", ),
                    "cfg": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 100.0, "step": 0.1, "tooltip": "Works with a list of floats too (one cfg float per step)"}),
                    "latent": ("LATENT", {"tooltip": "You can connect here the latent from TTM Latent Add, to pass reference video and ttm options"}),
                    "start_sampler_step": ("INT", {"default": 0, "min": 0, "max": 1000, "step": 1, "tooltip": "Start step of the whole sampling process. It will automatically skip the selected number of sigmas (starting from the first ones); if the sampler has a start_step option and you changed its value, set the same here"}),
                    },
                }

    RETURN_TYPES = ("GUIDER",)
    RETURN_NAMES = ("guider",)
    FUNCTION = "guide"
    CATEGORY = "Wan22 TimeToMove"

    def guide(cls, model, positive, negative, cfg, latent, start_sampler_step):
        guider = TTMGuider(model)
        guider.set_conds(positive, negative)
        guider.set_cfg(cfg)

        ttm_options = {}
        ttm_options["ttm_reference_latents"] = latent.get("ttm_reference_latents", None)
        ttm_options["ttm_start_step"] = latent["ttm_start_step"]
        ttm_options["ttm_end_step"] = latent["ttm_end_step"]
        ttm_options["latent_image"] = latent["samples"]
        ttm_options["motion_mask"] = latent["ttm_mask"]
        ttm_options["start_sampler_step"] = start_sampler_step
        guider.set_ttm_options(ttm_options)

        return (guider,)