import torch
import numpy as np
import math
import cv2
import sys, os
from matplotlib import image
from matplotlib import pyplot as plt
from models.pose_estimator.pose_estimator_model_setup import get_pose_estimation
from math import atan

# KEYPOINT_INDEXES = {
#     0: 'r ankle',
#     1: 'r knee',
#     2: 'r hip',
#     3: 'l hip',
#     4: 'l knee',
#     5: 'l ankle',
#     6: 'pelvis',
#     7: 'thorax',
#     8: 'upper neck',
#     9: 'head',
#     10: 'r wrist',
#     11: 'r elbow',
#     12: 'r shoulder',
#     13: 'l shoulder',
#     14: 'l elbow',
#     15: 'l wrist',
# }

def slope(x1, y1, x2, y2):
    if x1 == x2:
        return "undefined"
    return (y2-y1)/(x2-x1)

# Function to find the
# angle between two lines
def findAngle(M1, M2):
    vertical_line = False
    if M1 == "undefined":
        M1 = 0
        vertical_line = True
    if M2 == "undefined":
        M2 = 0
        vertical_line = True
    PI = 3.14159265
     
    # Store the tan value  of the angle
    angle = abs((M2 - M1) / (1 + M1 * M2))
 
    # Calculate tan inverse of the angle
    ret = atan(angle)
 
    # Convert the angle from
    # radian to degree
    val = (ret * 180) / PI
 
    # Print the result
    #print(round(val, 4))
    if vertical_line:
        return 90 - round(val,4)
    return (round(val, 4))

# np.array(pose_pred)[0][i]
def applyFeetApartError(filepath, pose_pred=None, diver_detector=None, pose_model=None):
    if pose_pred is None and filepath != "":
        diver_box, pose_pred = get_pose_estimation(filepath, diver_detector=diver_detector, pose_model=pose_model)
    if pose_pred is not None:
        pose_pred = np.array(pose_pred)[0]
        average_knee = [np.mean((pose_pred[4][0], pose_pred[1][0])), np.mean((pose_pred[4][1], pose_pred[1][1]))]
        vector1 = [pose_pred[5][0] - average_knee[0], pose_pred[5][1] - average_knee[1]]
        vector2 = [pose_pred[0][0] - average_knee[0], pose_pred[0][1] - average_knee[1]]
        unit_vector_1 = vector1 / np.linalg.norm(vector1)
        unit_vector_2 = vector2 / np.linalg.norm(vector2)
        dot_product = np.dot(unit_vector_1, unit_vector_2)
        angle = math.degrees(np.arccos(dot_product))
        # left_leg_slope = slope(np.array(pose_pred)[0][4][0], np.array(pose_pred)[0][4][1], np.array(pose_pred)[0][5][0], np.array(pose_pred)[0][5][1])
        # right_leg_slope = slope(np.array(pose_pred)[0][4][0], np.array(pose_pred)[0][4][1], np.array(pose_pred)[0][0][0], np.array(pose_pred)[0][0][1])
        # angle = findAngle(left_leg_slope, right_leg_slope)
        return angle
    else:
        # print('pose_pred is None')
        return None

#position: somersault or twist
def applyPositionTightnessError(filepath, pose_pred=None, diver_detector=None, pose_model=None):
    if pose_pred is None and filepath != "":
        diver_box, pose_pred = get_pose_estimation(filepath, diver_detector=diver_detector, pose_model=pose_model)
    if pose_pred is not None:
        pose_pred = np.array(pose_pred)[0]
        vector1 = [pose_pred[7][0] - pose_pred[2][0], pose_pred[7][1] - pose_pred[2][1]]
        
        vector2 = [pose_pred[1][0] - pose_pred[2][0], pose_pred[1][1] - pose_pred[2][1]]
        unit_vector_1 = vector1 / np.linalg.norm(vector1)
        unit_vector_2 = vector2 / np.linalg.norm(vector2)
        dot_product = np.dot(unit_vector_1, unit_vector_2)
        angle = math.degrees(np.arccos(dot_product))
        # upper_body = slope(np.array(pose_pred)[0][2][0], np.array(pose_pred)[0][2][1], np.array(pose_pred)[0][7][0], np.array(pose_pred)[0][7][1])
        # lower_body = slope(np.array(pose_pred)[0][2][0], np.array(pose_pred)[0][2][1], np.array(pose_pred)[0][1][0], np.array(pose_pred)[0][1][1])
        # angle = findAngle(upper_body, lower_body)
    # if position == 1:
    #     angle = 180 - angle
        return angle
    else:
        # print('pose_pred is None')
        return None