import numpy as np
import math

def distance_point_to_line_segment(px, py, x1, y1, x2, y2):
    # Calculate the squared distance from point (px, py) to the line segment [(x1, y1), (x2, y2)]
    def sqr_distance_point_to_segment():
        line_length_sq = (x2 - x1)**2 + (y2 - y1)**2
        if line_length_sq == 0:
            return (px - x1)**2 + (py - y1)**2
        t = max(0, min(1, ((px - x1) * (x2 - x1) + (py - y1) * (y2 - y1)) / line_length_sq))
        return ((px - (x1 + t * (x2 - x1)))**2 + (py - (y1 + t * (y2 - y1)))**2)

    # Calculate the closest point on the line segment to the given point (px, py)
    def closest_point_on_line_segment():
        line_length_sq = (x2 - x1)**2 + (y2 - y1)**2
        if line_length_sq == 0:
            return x1, y1
        t = max(0, min(1, ((px - x1) * (x2 - x1) + (py - y1) * (y2 - y1)) / line_length_sq))
        closest_x = x1 + t * (x2 - x1)
        closest_y = y1 + t * (y2 - y1)
        return closest_x, closest_y

    closest_point = closest_point_on_line_segment()
    distance = math.sqrt(sqr_distance_point_to_segment())

    return closest_point, distance

def min_distance_from_line_to_circle(line_start, line_end, circle_center, circle_radius):
    closest_point, distance = distance_point_to_line_segment(circle_center[0], circle_center[1],
                                                             line_start[0], line_start[1],
                                                             line_end[0], line_end[1])

    min_distance = max(0, distance - circle_radius)
    return min_distance


# def twist_counter(pose_pred, prev_pose_pred=None, in_petal=False, petal_count=0):
#     if pose_pred is None:
#         return petal_count, in_petal
#     min_dist = 0
#     # Users/lokamoto/Comprehensive_AQA/output/joint_plots/FINAWorldChampionships2019_Women10m_final_r1_0
#     pose_pred = pose_pred[0]
#     vector1 = [pose_pred[2][0] - pose_pred[3][0], 0-(pose_pred[2][1] - pose_pred[3][1])]
#     if prev_pose_pred is not None:
#         prev_pose_pred = prev_pose_pred[0]
#         prev_pose_pred = [prev_pose_pred[2][0] - prev_pose_pred[3][0], 0-(prev_pose_pred[2][1] - prev_pose_pred[3][1])]
#         # m = (vector1[1] - prev_pose_pred[1])/(vector1[0] - prev_pose_pred[0])
#         # b = prev_pose_pred[1] - m * prev_pose_pred[0]
#         # min_dist = np.abs(b)/np.sqrt(m**2+1)
#         min_dist = min_distance_from_line_to_circle(prev_pose_pred, vector1, (0, 0), 5)
    
#     if min_dist is not None and in_petal and np.linalg.norm(vector1) > 5 and min_dist == 0: #and np.linalg.norm(vector1) > 8
#         petal_count += 1
#         # print('leaving petal')
#         # print('going in new petal')
#     elif not in_petal and np.linalg.norm(vector1) > 5: #and min_dist > 3: #and np.linalg.norm(vector1) > 8
#         in_petal = True
#         # print('going in petal')
#     elif in_petal and np.linalg.norm(vector1) < 5:
#         in_petal = False
#         petal_count += 1
#         # print('leaving petal')
#     # print(vector)
#     return petal_count, in_petal

def twist_counter(pose_pred, prev_pose_pred=None, in_petal=False, petal_count=0):
    # if key[0].startswith('FINAWorldChampionships2019'):
    #     valid = 19
    #     outer = 3
    #     inner = 3
    # elif key[0].startswith('FINADivingWorldCup2021'):
    # valid = 19
    # outer = 5.5
    # inner = 5.5
    valid = 17
    outer = 10
    inner = 9
    if pose_pred is None:
        return petal_count, in_petal
    min_dist = 0
    pose_pred = pose_pred[0]
    vector1 = [pose_pred[2][0] - pose_pred[3][0], 0-(pose_pred[2][1] - pose_pred[3][1])]
    if prev_pose_pred is not None:
        prev_pose_pred = prev_pose_pred[0]
        prev_pose_pred = [prev_pose_pred[2][0] - prev_pose_pred[3][0], 0-(prev_pose_pred[2][1] - prev_pose_pred[3][1])]
        min_dist = min_distance_from_line_to_circle(prev_pose_pred, vector1, (0, 0), 0.5)
    if np.linalg.norm(vector1) > valid:
        return petal_count, in_petal
    if min_dist is not None and in_petal and np.linalg.norm(vector1) > outer and min_dist == 0: #and np.linalg.norm(vector1) > 8
        petal_count += 1
    elif not in_petal and np.linalg.norm(vector1) > outer: #and min_dist > 3: #and np.linalg.norm(vector1) > 8
        in_petal = True
    elif in_petal and np.linalg.norm(vector1) < inner:
        in_petal = False
        petal_count += 1
    return petal_count, in_petal

# def som_counter(pose_pred=None, prev_pose_pred=None, half_som_count=0, handstand=False):
#     if pose_pred is None:
#         return half_som_count
#     pose_pred = pose_pred[0]
#     vector1 = [pose_pred[7][0] - pose_pred[6][0], 0-(pose_pred[7][1] - pose_pred[6][1])] # flip y axis
#     if (not handstand and half_som_count % 2 == 0) or (handstand and half_som_count %2 == 1):  
#         vector2 = [0, -1]
#     else:
#         vector2 = [0, 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)
#     current_angle = math.degrees(np.arccos(dot_product))
    
#     if prev_pose_pred is not None:
#         prev_pose_pred = prev_pose_pred[0]
#         prev_vector = [prev_pose_pred[7][0] - prev_pose_pred[6][0], 0-(prev_pose_pred[7][1] - prev_pose_pred[6][1])] # flip y axis
#         prev_unit_vector = prev_vector / np.linalg.norm(prev_vector)
#         prev_angle_diff = math.degrees(np.arccos(np.dot(unit_vector_1, prev_unit_vector)))
#         # if prev_angle_diff > 120:
#         #     print('pose pred is probably off')
#         #     return half_som_count

#     # print("unit_vector_1:", unit_vector_1)
#     # print("looking for vector:", vector2)
#     if current_angle < 80:
#         half_som_count += 1
#     return half_som_count
    
def som_counter(pose_pred=None, prev_pose_pred=None, half_som_count=0, handstand=False):
    if pose_pred is None:
        return half_som_count, True
    pose_pred = pose_pred[0]
    vector1 = [pose_pred[7][0] - pose_pred[6][0], 0-(pose_pred[7][1] - pose_pred[6][1])] # flip y axis
    if (not handstand and half_som_count % 2 == 0) or (handstand and half_som_count %2 == 1):  
        vector2 = [0, -1]
    else:
        vector2 = [0, 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)
    current_angle = math.degrees(np.arccos(dot_product))
    if prev_pose_pred is not None:
        prev_pose_pred = prev_pose_pred[0]
        prev_vector = [prev_pose_pred[7][0] - prev_pose_pred[6][0], 0-(prev_pose_pred[7][1] - prev_pose_pred[6][1])] # flip y axis
        prev_unit_vector = prev_vector / np.linalg.norm(prev_vector)
        prev_angle_diff = math.degrees(np.arccos(np.dot(unit_vector_1, prev_unit_vector)))
        if prev_angle_diff > 115:
            return half_som_count, True

    # print("unit_vector_1:", unit_vector_1)
    # print("looking for vector:", vector2)
    if current_angle <= 80:
        half_som_count += 1
    return half_som_count, False