import os

import cv2
import easypose as ep
import numpy as np

# ==================================================================================================

filepath = os.path.dirname(os.path.realpath(__file__)) + "/"

# ==================================================================================================


def load_model():
    print("Loading mmpose model ...")

    model = ep.TopDown(
        "/RapidPoseTriangulation/extras/mmdeploy/exports/rtmpose-m_384x288_fp16_extra-steps.onnx",
        "SimCC",
        "/RapidPoseTriangulation/extras/mmdeploy/exports/rtmdet-nano_320x320_fp16_extra-steps.onnx",
        conf_threshold=0.3,
        iou_threshold=0.3,
        warmup=10,
    )

    print("Loaded mmpose model")
    return model


def load_wb_model():
    print("Loading mmpose whole body model ...")

    model = None

    print("Loaded mmpose model")
    return model


# ==================================================================================================


def get_2d_pose(model, imgs, num_joints=17):
    """See: https://mmpose.readthedocs.io/en/latest/user_guides/inference.html#basic-usage"""

    new_poses = []
    for i in range(len(imgs)):
        img = imgs[i]
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

        poses = []
        dets = model.predict(img)
        for pose in dets:
            pose = pose.keypoints
            pose = np.asarray(pose)

            scores = pose[:, 2].reshape(-1, 1)
            scores = np.clip(scores, 0, 1)
            pose = np.concatenate((pose[:, :2], scores), axis=-1)

            poses.append(pose)

        if len(poses) == 0:
            poses.append(np.zeros([num_joints, 3]))

        poses = np.array(poses)
        new_poses.append(poses)

    return new_poses