import copy

import cv2
import numpy as np

from skelda import utils_pose

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

core_joints = [
    "shoulder_middle",
    "hip_middle",
    "shoulder_left",
    "shoulder_right",
    "hip_left",
    "hip_right",
    "elbow_left",
    "elbow_right",
    "knee_left",
    "knee_right",
    "wrist_left",
    "wrist_right",
    "ankle_left",
    "ankle_right",
]

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


def undistort_points(points: np.ndarray, caminfo: dict):
    K = np.array(caminfo["K"], dtype=np.float32)
    DC = np.array(caminfo["DC"][0:5], dtype=np.float32)
    w = caminfo["width"]
    h = caminfo["height"]

    # Undistort camera matrix
    newK, _ = cv2.getOptimalNewCameraMatrix(K, DC, (w, h), 1, (w, h))
    caminfo = copy.deepcopy(caminfo)
    caminfo["K"] = newK.tolist()
    caminfo["DC"] = [0.0, 0.0, 0.0, 0.0, 0.0]

    # Undistort points
    pshape = points.shape
    points = np.reshape(points, (-1, 1, 2))
    points = cv2.undistortPoints(points, K, DC, P=newK)
    points = points.reshape(pshape)

    return points, caminfo


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


def get_camera_P(cam):
    """Calculate opencv-style projection matrix"""

    R = np.array(cam["R"])
    T = np.array(cam["T"])
    Tr = R @ (T * -1)
    P = cam["K"] @ np.hstack([R, Tr])
    return P


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


def calc_pair_score(pair, poses_2d, camparams, joint_names_2d):
    """Triangulates a pair of persons and scores them based on the reprojection error"""

    cam1 = camparams[pair[0][0]]
    cam2 = camparams[pair[0][1]]
    pose1 = np.array(poses_2d[pair[0][0]][pair[0][2]])
    pose2 = np.array(poses_2d[pair[0][1]][pair[0][3]])

    # Select core joints
    jids = [joint_names_2d.index(j) for j in core_joints]
    pose1 = pose1[jids]
    pose2 = pose2[jids]

    poses_3d, score = calc_pose_scored(pose1, pose2, cam1, cam2)
    return poses_3d, score


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


def calc_pose_scored(pose1, pose2, cam1, cam2):
    """Triangulates a pair of persons and scores them based on the reprojection error"""

    # Triangulate points
    points1 = pose1[:, 0:2].T
    points2 = pose2[:, 0:2].T
    P1 = get_camera_P(cam1)
    P2 = get_camera_P(cam2)
    points3d = cv2.triangulatePoints(P1, P2, points1, points2)
    points3d = points3d / points3d[3, :]
    points3d = points3d[0:3, :].T

    # Calculate reprojection error
    poses_3d = np.expand_dims(points3d, axis=0)
    poses_3d = np.concatenate([poses_3d, np.ones([1, poses_3d.shape[1], 1])], axis=-1)
    repro1, _ = utils_pose.project_poses(poses_3d, cam1)
    repro2, _ = utils_pose.project_poses(poses_3d, cam2)
    repro1 = repro1[0]
    repro2 = repro2[0]
    mask1 = pose1[:, 2] > 0.1
    mask2 = pose2[:, 2] > 0.1
    mask = mask1 & mask2
    error1 = np.linalg.norm(pose1[mask, 0:2] - repro1[mask, 0:2], axis=1)
    error2 = np.linalg.norm(pose2[mask, 0:2] - repro2[mask, 0:2], axis=1)

    # Convert errors to a score
    error1 = error1 / ((cam1["width"] + cam1["height"]) / 2)
    error2 = error2 / ((cam2["width"] + cam2["height"]) / 2)
    error = (error1 + error2) / 2
    scores = 1.0 / (1.0 + error * 10)
    score = np.mean(scores)

    # Add score to 3D pose
    full_scores = np.zeros([poses_3d.shape[1], 1])
    full_scores[mask] = np.expand_dims(scores, axis=-1)
    pose_3d = np.concatenate([points3d, full_scores], axis=-1)

    return pose_3d, score


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


def calc_grouping(all_pairs):
    """Groups pairs that share a person"""

    # Calculate the pose center for each pair
    for i in range(len(all_pairs)):
        pair = all_pairs[i]
        pose_3d = pair[2][0]
        mask = pose_3d[:, 2] > 0.1
        center = np.mean(pose_3d[mask, 0:3], axis=0)
        all_pairs[i] = all_pairs[i] + (center,)

    groups = []
    for i in range(len(all_pairs)):
        pair = all_pairs[i]

        # Create new group if non exists
        if len(groups) == 0:
            groups.append([pair])
            continue

        # Check if the pair belongs to an existing group
        matched = False
        for j in range(len(groups)):
            g0 = groups[j][0]
            center0 = g0[3]
            center1 = pair[3]
            if np.linalg.norm(center0 - center1) < 0.5:
                pose0 = g0[2][0]
                pose1 = pair[2][0]

                # Calculate the distance between the two poses
                mask0 = pose0[:, 3] > 0.1
                mask1 = pose1[:, 3] > 0.1
                mask = mask0 & mask1
                dists = np.linalg.norm(pose0[mask, 0:3] - pose1[mask, 0:3], axis=1)
                dist = np.mean(dists)
                if dist < 0.3:
                    groups[j].append(pair)
                    matched = True
                    break

        # Create new group if no match was found
        if not matched:
            groups.append([pair])

    return groups


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


def merge_group(group, poses_2d, camparams):
    """Merges a group of poses into a single pose"""

    # Calculate full 3D poses
    poses_3d = []
    for pair in group:
        cam1 = camparams[pair[0][0]]
        cam2 = camparams[pair[0][1]]
        pose1 = np.array(poses_2d[pair[0][0]][pair[0][2]])
        pose2 = np.array(poses_2d[pair[0][1]][pair[0][3]])

        pose_3d, _ = calc_pose_scored(pose1, pose2, cam1, cam2)
        poses_3d.append(pose_3d)

    # Merge poses
    pose_3d = np.mean(poses_3d, axis=0)

    return pose_3d


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


def get_3d_pose(poses_2d, camparams, joint_names_2d):

    # Undistort 2D points
    for i, cam in enumerate(camparams):
        poses = np.array(poses_2d[i])
        poses[:, :, 0:2], camparams[i] = undistort_points(poses[:, :, 0:2], cam)
        poses_2d[i] = poses.tolist()

    # Create pairs of persons
    num_persons = [len(p) for p in poses_2d]
    all_pairs = []
    for i in range(len(poses_2d)):
        poses = poses_2d[i]
        for j in range(i + 1, len(poses_2d)):
            poses2 = poses_2d[j]
            for k in range(len(poses)):
                for l in range(len(poses2)):
                    pid1 = sum(num_persons[:i]) + k
                    pid2 = sum(num_persons[:j]) + l
                    all_pairs.append(((i, j, k, l), (pid1, pid2)))

    # Calculate pair scores
    for i in range(len(all_pairs)):
        pair = all_pairs[i]
        pose_3d, score = calc_pair_score(pair, poses_2d, camparams, joint_names_2d)
        all_pairs[i] = all_pairs[i] + ((pose_3d, score),)

        # import draw_utils
        # poses3D = np.array([pose_3d])
        # _ = draw_utils.utils_view.show_poses3d(
        #     poses3D, core_joints, {}, camparams
        # )
        # draw_utils.utils_view.show_plots()

    # Drop pairs with low scores
    min_score = 0.9
    all_pairs = [p for p in all_pairs if p[2][1] > min_score]

    # Group pairs that share a person
    groups = calc_grouping(all_pairs)

    # Merge groups
    poses_3d = []
    for group in groups:
        pose_3d = merge_group(group, poses_2d, camparams)
        poses_3d.append(pose_3d)

    if len(poses_3d) > 0:
        poses3D = np.array(poses_3d)
    else:
        poses3D = np.zeros([1, len(joint_names_2d), 4])
    return poses3D