import numpy as np
from ultralytics.engine.results import Results


KEYPOINT_NAMES = ["nose","left_eye","right_eye","left_ear","right_ear","left_shoulder",
                "right_shoulder","left_elbow","right_elbow","left_wrist","right_wrist",
                "left_hip","right_hip","left_knee","right_knee","left_ankle","right_ankle"]


def get_keypoints(result: Results):
    keypoints = None
    for i, box in enumerate(result.boxes):
        if box.cls != 0.:  # Only consider the person class
            continue
        person_conf = box.conf.item()
        k = result.keypoints.data[i]
        x = k[:, 0].tolist()
        y = k[:, 1].tolist()
        score = k[:, 2]
        visible = (score > 0.5).tolist()
        # keypoints = {'x': x, 'y': y, 'visible': visible}
        keypoints = {key_name: (x_, y_) if v_ else None for key_name,x_,y_,v_ in zip(KEYPOINT_NAMES, x, y, visible)}
        break
    return keypoints


def calculate_angle(p1, p2, p3):
    v1 = np.array([p1[0] - p2[0], p1[1] - p2[1]])
    v2 = np.array([p3[0] - p2[0], p3[1] - p2[1]])
    angle_rad = np.arccos(np.dot(v1, v2) / (np.linalg.norm(v1) * np.linalg.norm(v2)))
    angle_deg = np.degrees(angle_rad)
    return angle_deg


def calculate_angle_to_horizontal(vector):
    angle_rad = np.arctan2(vector[1], vector[0])
    angle_deg = np.degrees(angle_rad)
    # Adjust the angle to be within -90 to +90 degrees
    if angle_deg > 90:
        angle_deg = 180 - angle_deg
    elif angle_deg < -90:
        angle_deg = -180 - angle_deg
    return -angle_deg


def extend_line(start, end, extend_factor=3):
    vector = np.array(end) - np.array(start)
    length = np.linalg.norm(vector)
    unit_vector = vector / np.linalg.norm(vector)
    new_point = end + unit_vector * length * extend_factor
    new_point = new_point.tolist()
    return (new_point[0], new_point[1])


def get_elbow_angles(keypoints: dict):
    left_elbow_angle = None
    right_elbow_angle = None
    if keypoints['left_shoulder'] and keypoints['left_elbow'] and keypoints['left_wrist']:
        left_elbow_angle = calculate_angle(keypoints['left_shoulder'], keypoints['left_elbow'], keypoints['left_wrist'])
    if keypoints['right_shoulder'] and keypoints['right_elbow'] and keypoints['right_wrist']:
        right_elbow_angle = calculate_angle(keypoints['right_shoulder'], keypoints['right_elbow'], keypoints['right_wrist'])
    return left_elbow_angle, right_elbow_angle


def get_eye_angles(keypoints: dict):
    left_eye_angle = None
    right_eye_angle = None
    if keypoints['left_ear'] and keypoints['left_eye']:
        left_vector = (keypoints['left_eye'][0] - keypoints['left_ear'][0], keypoints['left_eye'][1] - keypoints['left_ear'][1])
        left_eye_angle = calculate_angle_to_horizontal(left_vector)
    if keypoints['right_ear'] and keypoints['right_eye']:
        right_vector = (keypoints['right_eye'][0] - keypoints['right_ear'][0], keypoints['right_eye'][1] - keypoints['right_ear'][1])
        right_eye_angle = calculate_angle_to_horizontal(right_vector)
    return left_eye_angle, right_eye_angle