import numpy as np



class PoseMatching():
    """
    num_joints : 존재하는 joint 갯수
    pck_ratio_thresh : (0~1) 값이 낮으면 더욱 비슷해야 정답으로 인정.
    """
    def __init__(self,num_joints = 16, pck_ratio_thresh = 0.1):
        self.num_joints = num_joints
        self.pck_ratio_thresh = pck_ratio_thresh
        pass

    def _enclosing_rect(self,points):
        xs = points[0::2]
        ys = points[1::2]
        return np.array([np.amin(xs), np.amin(ys), np.amax(xs), np.amax(ys)])

    def rect_size(self,rect):
        return np.array([rect[2] - rect[0], rect[3] - rect[1]])

    """
    input
    prediction_joint : 예측한 좌표
    groundtruth_joint : 정답 좌표
    gt_present_joints : 정답 좌표에서 joint 가 존재하는지 안하는지

    output
    pck : 0~100. 100 이면 모두 정확히 일치.
    """
    def _eval_pck(self,prediction_joint, groundtruth_joint,gt_present_joints,gt_rect):

        pred_joints = np.zeros((self.num_joints, 2))
        gt_joints = np.zeros((self.num_joints, 2))

        if len(groundtruth_joint) == 0:
            raise Exception("gt joint is 0 .")

        gt_joints[:, 0] = groundtruth_joint[0::2]
        gt_joints[:, 1] = groundtruth_joint[1::2]

        pck_thresh = self.pck_ratio_thresh*np.amax(self.rect_size(gt_rect))

        pred_joints[:, 0] = prediction_joint[0::2]
        pred_joints[:, 1] = prediction_joint[1::2]

        dists = np.sqrt(np.sum((pred_joints - gt_joints)**2, axis=1))
        correct = dists <= pck_thresh

        num_all = np.sum(gt_present_joints, axis=0)

        num_correct = np.zeros((self.num_joints, ))
        for j_id in range(self.num_joints):
            num_correct[j_id] = np.sum(correct[gt_present_joints[j_id] == 1, j_id], axis=0)

        pck = np.sum(num_correct,axis=0)/num_all*100.0

        return pck


if __name__ == "__main__":

    num_joint = 3
    pck_ratio_thresh = 0.1
    match_mod = PoseMatching(num_joints=num_joint, pck_ratio_thresh=pck_ratio_thresh)
    # 템플릿 생성시 아래 한번 call
    gt_present_joints = [1,0,1]
    groundtruth_joint = [11,21,13,26,45,50]
    gt_box = match_mod._enclosing_rect(groundtruth_joint)


    # prediction 할때마다 아래처럼 matching 정확도 측정
    prediction_joint = [10,20,13,25,40,45]
    pck = match_mod._eval_pck(prediction_joint,groundtruth_joint,gt_present_joints, gt_box)

    print(pck)