# face_keypoints_filter.py
#
# DWPose / OpenPose の POSE_KEYPOINT から
# 顔パーツ（あご / 眉 / 目 / 鼻 / 口外周 / 口内周）を部分的に無効化するノード
#
# 前提:
#   pose_kps は comfyui_controlnet_aux DWPose(DWPreprocessor) が出す
#   OpenPose 互換 JSON 形式:
#   [
#     {
#       "version": 1.3,
#       "people": [
#         {
#           "pose_keypoints_2d": [...],
#           "face_keypoints_2d": [...],
#           "hand_left_keypoints_2d": [...],
#           "hand_right_keypoints_2d": [...]
#         }
#       ],
#       "canvas_width": 512,
#       "canvas_height": 512
#     },
#     ...
#   ]


class FaceKeypointsFilter:
    """
    POSE_KEYPOINT 内の face_keypoints_2d から
    特定の顔パーツの信頼度(score)を 0 にして無効化するノード。
    DWPose / OpenPose の JSON 構造 (frame["people"][0]["face_keypoints_2d"]) に対応。
    """

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "pose_kps": ("POSE_KEYPOINT",),
                # True のパーツは「無効化」される（score=0）
                "remove_jaw": ("BOOLEAN", {"default": False}),
                "remove_eyebrows": ("BOOLEAN", {"default": False}),
                "remove_eyes": ("BOOLEAN", {"default": False}),
                "remove_nose": ("BOOLEAN", {"default": False}),
                # ★ 口を外周・内周で分ける
                "remove_mouth_outer": ("BOOLEAN", {"default": True}),
                "remove_mouth_inner": ("BOOLEAN", {"default": False}),
            }
        }

    RETURN_TYPES = ("POSE_KEYPOINT",)
    RETURN_NAMES = ("pose_kps",)
    FUNCTION = "filter_keypoints"
    CATEGORY = "Pose/Face Utils"

    # ------------------ internal helpers ------------------ #

    def _zero_indices(self, flat_list, indices):
        """
        flat_list: [x0,y0,c0, x1,y1,c1, ...] の配列
        indices: 無効化したい「点 index」のリスト (0〜67)
        該当 index の (x,y,score) を 0 にする
        """
        if flat_list is None:
            return flat_list

        if len(flat_list) % 3 != 0:
            # 想定外の形式はそのまま返す
            return flat_list

        num_points = len(flat_list) // 3
        if num_points < 68:
            # 顔68点が入っていない場合は何もしない
            return flat_list

        out = list(flat_list)
        for idx in indices:
            if idx < 0 or idx >= num_points:
                continue
            base = idx * 3
            out[base] = 0.0
            out[base + 1] = 0.0
            out[base + 2] = 0.0

        return out

    def _get_face_list_from_frame(self, frame):
        """
        frame から face_keypoints_2d のフラット配列を取得する。
        DWPose 形式 (frame["people"][0]["face_keypoints_2d"]) と
        旧形式 (frame["face_keypoints_2d"]) の両方に対応。
        """
        if not isinstance(frame, dict):
            return None, None, None

        # DWPose / OpenPose 標準 JSON:
        # frame["people"][0]["face_keypoints_2d"]
        if "people" in frame and isinstance(frame["people"], list) and frame["people"]:
            person = frame["people"][0]
            if isinstance(person, dict) and "face_keypoints_2d" in person:
                return frame, person, person["face_keypoints_2d"]

        # 互換用：トップレベルに face_keypoints_2d がいる場合
        if "face_keypoints_2d" in frame:
            return frame, None, frame["face_keypoints_2d"]

        return None, None, None

    # ------------------ main function ------------------ #

    def filter_keypoints(
        self,
        pose_kps,
        remove_jaw,
        remove_eyebrows,
        remove_eyes,
        remove_nose,
        remove_mouth_outer,
        remove_mouth_inner,
    ):
        """
        pose_kps: DWPose / OpenPose 由来の POSE_KEYPOINT
        """

        # OpenPose 68point face のインデックス
        jaw_idx = list(range(0, 17))        # あご輪郭
        eyebrow_idx = list(range(17, 27))   # 眉
        nose_idx = list(range(27, 36))      # 鼻
        eye_idx = list(range(36, 48))       # 両目

        # ★ 口を外周・内周に分割
        mouth_outer_idx = list(range(48, 60))  # 48〜59: 唇の輪郭
        mouth_inner_idx = list(range(60, 68))  # 60〜67: 口の開閉ライン

        if not isinstance(pose_kps, list):
            # 想定：pose_kps は frame のリスト
            return (pose_kps,)

        new_pose_kps = []

        for frame in pose_kps:
            # フレームを浅いコピー
            if isinstance(frame, dict):
                frame_copy = dict(frame)
            else:
                new_pose_kps.append(frame)
                continue

            frame_ref, person_ref, face_flat = self._get_face_list_from_frame(frame_copy)

            if face_flat is None:
                # 顔情報が無ければそのまま
                new_pose_kps.append(frame_copy)
                continue

            indices_to_zero = []
            if remove_jaw:
                indices_to_zero += jaw_idx
            if remove_eyebrows:
                indices_to_zero += eyebrow_idx
            if remove_eyes:
                indices_to_zero += eye_idx
            if remove_nose:
                indices_to_zero += nose_idx
            if remove_mouth_outer:
                indices_to_zero += mouth_outer_idx
            if remove_mouth_inner:
                indices_to_zero += mouth_inner_idx

            if indices_to_zero:
                face_flat_new = self._zero_indices(list(face_flat), indices_to_zero)

                # person 内か frame 直下かで代入先を変える
                if person_ref is not None:
                    person_copy = dict(person_ref)
                    person_copy["face_keypoints_2d"] = face_flat_new
                    people_list = list(frame_copy.get("people", []))
                    if people_list:
                        people_list[0] = person_copy
                        frame_copy["people"] = people_list
                else:
                    frame_copy["face_keypoints_2d"] = face_flat_new

            new_pose_kps.append(frame_copy)

        return (new_pose_kps,)


NODE_CLASS_MAPPINGS = {
    "FaceKeypointsFilter": FaceKeypointsFilter,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "FaceKeypointsFilter": "Face Keypoints Filter (OpenPose)",
}