import cv2
import numpy as np

import torch
from mmpose.apis import inference_topdown, init_model
from mmpose.utils import register_all_modules
register_all_modules()

# coco keypoints: 
# https://github.com/open-mmlab/mmpose/blob/master/mmpose/datasets/datasets/top_down/topdown_coco_dataset.py#L28
model = init_model('mmpose/td-hm_hrnet-w48_8xb32-210e_coco-256x192.py', 'mmpose/td-hm_hrnet-w48_8xb32-210e_coco-256x192-0e67c616_20220913.pth', device=torch.cuda.current_device() if torch.cuda.is_available() else 'cpu')

def save_image(img, img_path):
    # Convert PIL image to OpenCV image
    img = cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)
    # Save OpenCV image
    cv2.imwrite(img_path, img)

def predict_poses(front_img, front_img_path, side_img, side_img_path):
    save_image(front_img, front_img_path)
    save_image(side_img, side_img_path)

    mmpose_result = mmpose_coco(front_img_path, side_img_path)
    front_result = mmpose_result[0]
    side_result = mmpose_result[1]

    front_keypoints = front_result[0].pred_instances['keypoints'][0]
    front_keypoints_data = {
        'keypoints': front_keypoints.tolist(),
        'keypoint_names': [
            'nose',
            'left_eye',
            'right_eye',
            'left_ear',
            'right_ear',
            'left_shoulder',
            'right_shoulder',
            'left_elbow',
            'right_elbow',
            'left_wrist',
            'right_wrist',
            'left_hip',
            'right_hip',
            'left_knee',
            'right_knee',
            'left_ankle',
            'right_ankle'
        ]
    }
    
    side_keypoints = side_result[0].pred_instances['keypoints'][0]
    side_keypoints_data = {
        'keypoints': side_keypoints.tolist(),
        'keypoint_names': [
            'nose',
            'left_eye',
            'right_eye',
            'left_ear',
            'right_ear',
            'left_shoulder',
            'right_shoulder',
            'left_elbow',
            'right_elbow',
            'left_wrist',
            'right_wrist',
            'left_hip',
            'right_hip',
            'left_knee',
            'right_knee',
            'left_ankle',
            'right_ankle'
        ]
    }

    return ((front_img, front_keypoints_data), (side_img, side_keypoints_data))

def mmpose_coco(front_img_path,
                side_img_path):

    front_results = inference_topdown(model, front_img_path)
    side_results = inference_topdown(model, side_img_path)

    return (front_results, side_results)