import os
import json
import numpy as np
from os import path


def match_pose_to_bbox(speaker_info_path, pose_json_dir, batch_index, count_keypoints_in_bbox_fn):
    # Load speaker info
    with open(speaker_info_path, 'r') as f:
        infos = json.load(f)

    # Select pose files
    pose_name_list = sorted([
        pose_name for pose_name in os.listdir(pose_json_dir)
        if pose_name.endswith('.json')
    ])
    pose_name_list = [pose_name_list[i] for i in batch_index]

    frame_pose_list = []
    for pose_name in pose_name_list:  # iterate over frames
        with open(path.join(pose_json_dir, pose_name), 'r') as f:
            pose_list = json.load(f)['instance_info']

        pose_in_order = []
        for info in infos:  # iterate over persons
            instance_bbox = info['instance_bbox']
            within_bbox_count = [
                count_keypoints_in_bbox_fn(pose, instance_bbox) for pose in pose_list
            ]
            match_index = np.argmax(within_bbox_count)
            pose_in_order.append(pose_list[match_index]['keypoints'])

        frame_pose_list.append(pose_in_order)

    return np.array(frame_pose_list).transpose(1, 0, 2, 3)  # (human, frame, 133, 2)
