'''
Example:

import cv2

from ComfyUI_PromptChain.vendor.rtmlib import Custom, draw_skeleton

device = 'cuda'
backend = 'onnxruntime'  # opencv, onnxruntime

cap = cv2.VideoCapture('./demo.mp4')

openpose_skeleton = False  # True for openpose-style, False for mmpose-style

# Example: RTMPose BodyWithFeet in balanced mode
custom = Custom(to_openpose=openpose_skeleton,
                det_class='YOLOX',
                det='https://download.openmmlab.com/mmpose/v1/projects/rtmposev1/onnx_sdk/yolox_m_8xb8-300e_humanart-c2c7a14a.zip', # noqa
                det_input_size=(640, 640),
                pose_class='RTMPose',
                pose='https://download.openmmlab.com/mmpose/v1/projects/rtmposev1/onnx_sdk/rtmpose-m_simcc-body7_pt-body7-halpe26_700e-256x192-4d3e73dd_20230605.zip', # noqa
                pose_input_size=(192, 256),
                backend=backend,
                device=device)

# # Example: ViTPose BodyWithFeet in balanced mode
# custom = Custom(to_openpose=openpose_skeleton,
#                 det_class='YOLOX',
#                 det='https://download.openmmlab.com/mmpose/v1/projects/rtmposev1/onnx_sdk/yolox_m_8xb8-300e_humanart-c2c7a14a.zip', # noqa
#                 det_input_size=(640, 640),
#                 pose_class='ViTPose',
#                 pose='https://huggingface.co/JunkyByte/easy_ViTPose/resolve/main/onnx/coco_25/vitpose-l-coco_25.onnx', # noqa
#                 pose_input_size=(192, 256),
#                 backend=backend,
#                 device=device)

# # Example: RTMO in balanced mode
# custom = Custom(to_openpose=openpose_skeleton,
#                 pose_class='RTMO',
#                 pose='https://download.openmmlab.com/mmpose/v1/projects/rtmo/onnx_sdk/rtmo-m_16xb16-600e_body7-640x640-39e78cc4_20231211.zip', # noqa
#                 pose_input_size=(640,640),
#                 backend=backend,
#                 device=device)

# # Example: Hand in lightweight mode
# custom = Custom(to_openpose=openpose_skeleton,
#                 det_class='RTMDet',
#                 det='https://download.openmmlab.com/mmpose/v1/projects/rtmposev1/onnx_sdk/rtmdet_nano_8xb32-300e_hand-267f9c8f.zip',
#                 det_input_size=(320,320),
#                 pose_class='RTMPose',
#                 pose='https://download.openmmlab.com/mmpose/v1/projects/rtmposev1/onnx_sdk/rtmpose-m_simcc-hand5_pt-aic-coco_210e-256x256-74fb594_20230320.zip',
#                 pose_input_size=(256, 256),
#                 backend=backend,
#                 device=device)

# # Example: Human and animal
# custom = Custom(to_openpose=openpose_skeleton,
#                 det_class='YOLOX',
#                 det_mode='multiclass', # or det_categories=[0,23] for example,
#                 det='https://github.com/Megvii-BaseDetection/YOLOX/releases/download/0.1.1rc0/yolox_s.onnx',
#                 det_input_size=(640,640),
#                 pose_class='ViTPose',
#                 pose='https://huggingface.co/JunkyByte/easy_ViTPose/resolve/main/onnx/apt36k/vitpose-b-apt36k.onnx',
#                 pose_input_size=(192, 256),
#                 backend=backend,
#                 device=device)

frame_idx = 0

while cap.isOpened():
    success, frame = cap.read()
    frame_idx += 1

    if not success:
        break

    keypoints, scores = custom(frame)

    img_show = frame.copy()

    img_show = draw_skeleton(img_show,
                             keypoints,
                             scores,
                             openpose_skeleton=openpose_skeleton,
                             kpt_thr=0.43)

    img_show = cv2.resize(img_show, (960, 540))
    cv2.imshow('img', img_show)
    cv2.waitKey(10)

'''
import importlib

import numpy as np


def _rtmlib_root():
    # Vendored: resolve the package root by its own relative path (…vendor.rtmlib),
    # never the top-level name 'rtmlib', so a pip-installed rtmlib can't shadow us
    # and import never fails when this slice lives under a renamed parent package.
    return importlib.import_module(__name__.rsplit('.tools.', 1)[0])


class Custom:

    def __init__(self,
                 det_class: str = None,
                 det: str = None,
                 det_input_size: tuple = (640, 640),
                 det_mode: str = 'human',
                 det_categories: list = None,
                 pose_class: str = None,
                 pose: str = None,
                 pose_input_size: tuple = (192, 256),
                 mode: str = None,
                 to_openpose: bool = False,
                 backend: str = 'onnxruntime',
                 device: str = 'cpu'):

        if det_class is not None:
            try:
                if det_categories is not None:
                    det_mode = 'multiclass'

                det_class = getattr(_rtmlib_root(), det_class)
                self.det_model = det_class(det,
                                           model_input_size=det_input_size,
                                           mode=det_mode,
                                           backend=backend,
                                           device=device)
                self.det_mode = det_mode
                self.det_categories = det_categories
                self.one_stage = False

            except ImportError:
                raise ImportError(f'{det_class} is not supported by rtmlib.')
        else:
            self.one_stage = True

        if pose_class is not None:
            try:
                pose_class = getattr(_rtmlib_root(), pose_class)
                self.pose_model = pose_class(pose,
                                             model_input_size=pose_input_size,
                                             to_openpose=to_openpose,
                                             backend=backend,
                                             device=device)
            except ImportError:
                raise ImportError(f'{pose_class} is not supported by rtmlib.')

    def __call__(self, image: np.ndarray):
        if self.one_stage:
            keypoints, scores = self.pose_model(image)
        else:
            if self.det_categories or self.det_mode == 'multiclass':
                if self.det_categories:
                    bboxes, classes = self.det_model(image)
                    bboxes = [
                        bbox for bbox, cls in zip(bboxes, classes)
                        if cls in self.det_categories
                    ]
                else:
                    bboxes, _ = self.det_model(image)
            else:
                bboxes = self.det_model(image)
            keypoints, scores = self.pose_model(image, bboxes=bboxes)

        return keypoints, scores