import torch
import comfy.model_management
import comfy.utils
import node_helpers


class WanI2VContinuity_Dolphin:
    """
    Wan2.2 I2V 노드 — 클립 이어붙이기(연속성)용.

    기존 접근(WanImageToVideoSVIPro)의 문제: anchor 이미지를 따로 인코딩한 latent와,
    이전 클립 latent의 뒷부분을 그냥 잘라낸 조각(motion_latent)을 이어붙였다. WAN VAE는
    인과적(causal) 시간 압축이라 각 latent 프레임이 "그 시점까지의 시퀀스 문맥"을 전제로
    인코딩되는데, 서로 다른 문맥에서 나온 latent 조각을 이어붙이면 통계적으로 불연속이라
    화질이 튀고 동작 연속성도 못 살렸다.

    이 노드는 ComfyUI 코어 `WanImageToVideo`의 검증된 방식을 따른다:
    이전 클립에서 디코딩한 마지막 N프레임(픽셀)을 start_image로 받아, 회색(0.5) filler와
    함께 하나의 시퀀스로 만들어 **한 번의 통짜 VAE 인코딩**을 한다. 그래서 known 구간과
    생성 구간의 latent 문맥이 자연스럽게 이어진다. 마스크도 프레임수→latent수 변환을
    정확히 해서 known 프레임 전체를 고정한다.

    여기에 PainterI2V의 motion_amplitude(슬로우모션 방지)를 이식했다. 단 known 프레임은
    건드리지 않고, 뒤에 붙는 filler(생성 대상) latent만 마지막 known latent 대비 차이를
    증폭한다. motion_amplitude=1.0이면 코어 WanImageToVideo와 동일하게 동작한다.
    """

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "positive": ("CONDITIONING",),
                "negative": ("CONDITIONING",),
                "vae": ("VAE",),
                "width": ("INT", {"default": 832, "min": 16, "max": 4096, "step": 16}),
                "height": ("INT", {"default": 480, "min": 16, "max": 4096, "step": 16}),
                "length": ("INT", {"default": 81, "min": 1, "max": 9999, "step": 4}),
                "batch_size": ("INT", {"default": 1, "min": 1, "max": 4096}),
                "motion_amplitude": ("FLOAT", {"default": 1.0, "min": 1.0, "max": 2.0, "step": 0.05,
                                                 "tooltip": "1.0=코어 WanImageToVideo와 동일. >1.0이면 filler 구간에 슬로우모션 방지 증폭."}),
            },
            "optional": {
                "clip_vision_output": ("CLIP_VISION_OUTPUT",),
                "start_image": ("IMAGE", {"tooltip": "known 프레임(들). 이전 클립의 디코딩된 마지막 N프레임을 넣으면 연속성이 생김. 1장이면 일반 I2V."}),
            },
        }

    RETURN_TYPES = ("CONDITIONING", "CONDITIONING", "LATENT")
    RETURN_NAMES = ("positive", "negative", "latent")
    FUNCTION = "encode"
    CATEGORY = "Dolphin Node/Wan Video"

    def encode(self, positive, negative, vae, width, height, length, batch_size,
               motion_amplitude=1.0, clip_vision_output=None, start_image=None):
        latent = torch.zeros(
            [batch_size, 16, ((length - 1) // 4) + 1, height // 8, width // 8],
            device=comfy.model_management.intermediate_device(),
        )

        if start_image is not None:
            start_image = comfy.utils.common_upscale(
                start_image[:length].movedim(-1, 1), width, height, "bilinear", "center"
            ).movedim(1, -1)
            n_known_frames = start_image.shape[0]

            # known 프레임 뒤를 회색(0.5) filler로 채운 뒤 통짜로 인코딩 (코어 방식)
            image = torch.ones(
                (length, height, width, start_image.shape[-1]),
                device=start_image.device, dtype=start_image.dtype,
            ) * 0.5
            image[:n_known_frames] = start_image

            concat_latent_image = vae.encode(image[:, :, :, :3])

            # 프레임수 -> latent수 (WAN VAE 4x 시간 압축)
            n_known_latents = ((n_known_frames - 1) // 4) + 1

            # motion_amplitude: filler(생성 대상) latent만 증폭. known 프레임은 그대로.
            if motion_amplitude > 1.0 and concat_latent_image.shape[2] > n_known_latents:
                base_latent = concat_latent_image[:, :, n_known_latents - 1:n_known_latents]
                filler = concat_latent_image[:, :, n_known_latents:]
                diff = filler - base_latent
                diff_mean = diff.mean(dim=(1, 3, 4), keepdim=True)
                diff_centered = diff - diff_mean
                scaled = base_latent + diff_centered * motion_amplitude + diff_mean
                scaled = torch.clamp(scaled, -6, 6)
                concat_latent_image = torch.cat(
                    [concat_latent_image[:, :, :n_known_latents], scaled], dim=2
                )

            mask = torch.ones(
                (1, 1, latent.shape[2], concat_latent_image.shape[-2], concat_latent_image.shape[-1]),
                device=start_image.device, dtype=start_image.dtype,
            )
            mask[:, :, :n_known_latents] = 0.0

            positive = node_helpers.conditioning_set_values(
                positive, {"concat_latent_image": concat_latent_image, "concat_mask": mask}
            )
            negative = node_helpers.conditioning_set_values(
                negative, {"concat_latent_image": concat_latent_image, "concat_mask": mask}
            )

        if clip_vision_output is not None:
            positive = node_helpers.conditioning_set_values(positive, {"clip_vision_output": clip_vision_output})
            negative = node_helpers.conditioning_set_values(negative, {"clip_vision_output": clip_vision_output})

        return (positive, negative, {"samples": latent})