# pose_keypoints_to_image.py
#
# POSE_KEYPOINT (DWPose / OpenPose 形式) から
# ポーズ画像（IMAGE）を描き直すノードです。

import numpy as np
from PIL import Image, ImageDraw
import torch


class PoseKeypointsToImage:
    """
    POSE_KEYPOINT からポーズ画像を生成するノード。
    DWPose の JSON (frame["people"][0]["pose_keypoints_2d"], etc.) に対応。
    FaceKeypointsFilter で口を消した後の pose_kps などを渡すことで、
    「口なし」の OpenPose 風画像を作ることができる。
    """

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "pose_kps": ("POSE_KEYPOINT",),
                "width": (
                    "INT",
                    {"default": 512, "min": 64, "max": 4096, "step": 8},
                ),
                "height": (
                    "INT",
                    {"default": 512, "min": 64, "max": 4096, "step": 8},
                ),
                "point_radius": (
                    "INT",
                    {"default": 2, "min": 1, "max": 10, "step": 1},
                ),
                "draw_body": ("BOOLEAN", {"default": True}),
                "draw_face": ("BOOLEAN", {"default": True}),
                "draw_hands": ("BOOLEAN", {"default": True}),
            }
        }

    RETURN_TYPES = ("IMAGE",)
    RETURN_NAMES = ("image",)
    FUNCTION = "render"
    CATEGORY = "Pose/Face Utils"

    # ------------------ helpers ------------------ #

    def _draw_points(self, draw, flat_list, radius, color=(255, 255, 255)):
        """
        flat_list: [x0, y0, s0, x1, y1, s1, ...]
        score(s) > 0 の点だけを描画する
        """
        if flat_list is None:
            return
        if len(flat_list) % 3 != 0:
            return

        for i in range(0, len(flat_list), 3):
            x, y, score = flat_list[i : i + 3]
            if score is None:
                continue
            if score <= 0:
                continue
            if np.isnan(x) or np.isnan(y):
                continue
            x = float(x)
            y = float(y)
            draw.ellipse(
                (x - radius, y - radius, x + radius, y + radius),
                fill=color,
            )

    def _get_people_list(self, frame):
        """
        frame["people"] があればそれを返し、
        無ければ [frame] を 1人分として扱う。
        """
        if not isinstance(frame, dict):
            return []

        if "people" in frame and isinstance(frame["people"], list) and frame["people"]:
            return frame["people"]

        # 旧形式: frame 自体に keypoints がいる想定
        return [frame]

    # ------------------ main ------------------ #

    def render(
        self,
        pose_kps,
        width,
        height,
        point_radius,
        draw_body,
        draw_face,
        draw_hands,
    ):
        """
        pose_kps: DWPose Estimator などからの POSE_KEYPOINT
        """

        # pose_kps は通常「各フレーム dict のリスト」
        if isinstance(pose_kps, dict):
            frames = [pose_kps]
        elif isinstance(pose_kps, list):
            frames = pose_kps
        else:
            blank_np = np.zeros((1, height, width, 3), dtype=np.float32)
            blank = torch.from_numpy(blank_np)
            return (blank,)

        # 可能なら 1フレーム目の canvas_width/height で上書き
        if frames and isinstance(frames[0], dict):
            cw = frames[0].get("canvas_width", width)
            ch = frames[0].get("canvas_height", height)
            try:
                width = int(cw)
                height = int(ch)
            except Exception:
                pass

        images = []

        for frame in frames:
            img = Image.new("RGB", (width, height), (0, 0, 0))
            draw = ImageDraw.Draw(img)

            if not isinstance(frame, dict):
                images.append(np.array(img).astype(np.float32) / 255.0)
                continue

            people = self._get_people_list(frame)

            for person in people:
                if not isinstance(person, dict):
                    continue

                if draw_body and "pose_keypoints_2d" in person:
                    self._draw_points(draw, person["pose_keypoints_2d"], point_radius)

                if draw_face and "face_keypoints_2d" in person:
                    self._draw_points(draw, person["face_keypoints_2d"], point_radius)

                if draw_hands:
                    if "hand_left_keypoints_2d" in person:
                        self._draw_points(
                            draw, person["hand_left_keypoints_2d"], point_radius
                        )
                    if "hand_right_keypoints_2d" in person:
                        self._draw_points(
                            draw, person["hand_right_keypoints_2d"], point_radius
                        )

            img_arr = np.array(img).astype(np.float32) / 255.0  # H,W,3
            images.append(img_arr)

        if not images:
            blank_np = np.zeros((1, height, width, 3), dtype=np.float32)
            blank = torch.from_numpy(blank_np)
            return (blank,)

        batch_np = np.stack(images, axis=0)   # (B,H,W,3)
        batch = torch.from_numpy(batch_np)    # torch.Tensor に変換

        return (batch,)


NODE_CLASS_MAPPINGS = {
    "PoseKeypointsToImage": PoseKeypointsToImage,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "PoseKeypointsToImage": "Pose Keypoints → Image",
}