import os
import torch
import numpy as np
from ..utils import log

from accelerate import init_empty_weights
from accelerate.utils import set_module_tensor_to_device

import comfy.model_management as mm
from comfy.utils import load_torch_file, ProgressBar
import folder_paths

script_directory = os.path.dirname(os.path.abspath(__file__))
device = mm.get_torch_device()
offload_device = mm.unet_offload_device()


class WanVideoAddSteadyDancerEmbeds:
    @classmethod
    def INPUT_TYPES(s):
        return {"required": {
                    "embeds": ("WANVIDIMAGE_EMBEDS",),
                    "pose_latents_positive": ("LATENT",),
                    "pose_strength_spatial": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 100.0, "step": 0.01, "tooltip": "Strength of the pose embedding"}),
                    "pose_strength_temporal": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 100.0, "step": 0.01, "tooltip": "Strength of the pose embedding"}),
                    "start_percent": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "Start percentage of the embedding application"}),
                    "end_percent": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "End percentage of the embedding application"}),
                },
                "optional": {
                    "pose_latents_negative": ("LATENT",),
                    "clip_vision_embeds": ("WANVIDIMAGE_CLIPEMBEDS",),
                    
                }
        }

    RETURN_TYPES = ("WANVIDIMAGE_EMBEDS",)
    RETURN_NAMES = ("image_embeds",)
    FUNCTION = "add"
    CATEGORY = "WanVideoWrapper"

    def add(self, embeds, pose_latents_positive, pose_strength_spatial, pose_strength_temporal, start_percent=0.0, end_percent=1.0, pose_latents_negative=None, clip_vision_embeds=None):
        sdancer_embeds = {
            "cond_pos": pose_latents_positive["samples"][0],
            "cond_neg": pose_latents_negative["samples"][0] if pose_latents_negative else None,
            "pose_strength_spatial": pose_strength_spatial,
            "pose_strength_temporal": pose_strength_temporal,
            "start_percent": start_percent,
            "end_percent": end_percent,
            "clip_fea": clip_vision_embeds,
        }

        updated = dict(embeds)
        updated["sdancer_embeds"] = sdancer_embeds
        return (updated,)


NODE_CLASS_MAPPINGS = {
    "WanVideoAddSteadyDancerEmbeds": WanVideoAddSteadyDancerEmbeds,
    }
NODE_DISPLAY_NAME_MAPPINGS = {
    "WanVideoAddSteadyDancerEmbeds": "WanVideo Add SteadyDancer Embeds",
    }