import cv2 as cv
import pandas as pd
import math
import numpy as np
from tqdm.notebook import tqdm


def encode(ann):
    keypoints_2d = pd.DataFrame(np.array(ann['keypoints']).reshape(-1, 3),
                                columns=['x','y','conf'], index=range(1, 25))
    bbox = np.array(ann['bbox'])
    score = ann['score']
    return keypoints_2d, bbox, score


def fit(imageSize, keypoints_2d, keypoints_3d, focus=1):

    objectPoints = keypoints_3d.loc[keypoints_2d.index].values
    imagePoints = keypoints_2d[['x', 'y']].values.astype('float')

    n = len(imagePoints)
    fx = fy = focus*np.hypot(*imageSize)
    cx = imageSize[0]/2
    cy = imageSize[1]/2
    distCoeffs = np.zeros(4, np.float32)
    if n < 6:
        raise ValueError('Number of keypoints must be > 5')

    cameraMatrix = np.float32([[fx,0, cx],
                               [0, fy,cy],
                               [0, 0, 1]])

    _, rvecs, tvecs = cv.solvePnP(objectPoints, imagePoints, cameraMatrix, distCoeffs, flags=cv.SOLVEPNP_ITERATIVE )
    return rvecs, tvecs, cameraMatrix, distCoeffs


def rvec2euler(rvec):
    rvec_matrix = cv.Rodrigues(rvec)[0]
    proj_matrix = np.hstack((rvec_matrix, np.zeros_like(rvec)))
    euler_angles = cv.decomposeProjectionMatrix(proj_matrix)[6]
    return euler_angles.flatten()

# Checks if a matrix is a valid rotation matrix.
def isRotationMatrix(R) :
    Rt = np.transpose(R)
    shouldBeIdentity = np.dot(Rt, R)
    I = np.identity(3, dtype = R.dtype)
    n = np.linalg.norm(I - shouldBeIdentity)
    return n < 1e-6

# Calculates rotation matrix to euler angles
# The result is the same as MATLAB except the order
# of the euler angles ( x and z are swapped ).
def rotationMatrixToEulerAngles(R) :

    assert(isRotationMatrix(R))

    sy = math.sqrt(R[0,0] * R[0,0] +  R[1,0] * R[1,0])

    singular = sy < 1e-6

    if  not singular :
        x = math.atan2(R[2,1] , R[2,2])
        y = math.atan2(-R[2,0], sy)
        z = math.atan2(R[1,0], R[0,0])
    else :
        x = math.atan2(-R[1,2], R[1,1])
        y = math.atan2(-R[2,0], sy)
        z = 0

    return np.array([x, y, z])

def rvec2euler_Fridman(rvec):
    R = cv.Rodrigues(rvec)[0]
    euler_angles = np.rad2deg(rotationMatrixToEulerAngles(R))
    return euler_angles


def pose_estimation(img_json, keypoints_3d, ImageSize=(1920, 1080), treshhold=0.01, focus=1, fridman=False):
    poses = []
    for i, ann in enumerate(img_json):
        keypoints_2d, bbox, score = encode(ann)

        conf_keypoints_2d = keypoints_2d[keypoints_2d['conf']>treshhold]
        if len(conf_keypoints_2d)>5:
            rvec, tvec, camMatrx, dist = fit(ImageSize, conf_keypoints_2d, keypoints_3d, focus=focus)
            if not fridman:
                euler_angles = rvec2euler(rvec)
            else:
                euler_angles = rvec2euler_Fridman(rvec) #alternative

            pose = {'Euler angles':[-euler_angles[2], -euler_angles[0]+90, -euler_angles[1]],
                    'xyz coords':[tvec[2][0],tvec[0][0],-tvec[1][0]]}
            poses.append(pose)
    return poses


def image_pose_estimation(img_json, img_name, *args, **kvargs):
    poses = pose_estimation(img_json, *args, **kvargs)
    UE4_json = {'data':{'filename':img_name, 'poses':poses}}
    return UE4_json


def video_pose_estimation(video_json, video_name, *args, **kvargs):
    video_poses = []
    for frame in tqdm(video_json):
        num = frame['frame']
        frame_poses = pose_estimation(frame['predictions'], *args, **kvargs)
        video_poses.append({'frame':num, 'poses':frame_poses})
    UE4_json_video = {'filename':video_name, 'data':video_poses}
    return UE4_json_video