import numpy as np
from scipy.spatial.distance import euclidean
from fastdtw import fastdtw

# Hằng số xác định độ lệch tối đa cho phép để coi là 'khớp' khi tìm frame bắt đầu
START_DETECTION_THRESHOLD = 15.0

def calculate_angle(a, b, c):
    """
    Tính góc 2D giữa 3 điểm. b là đỉnh của góc.
    Góc trả về đơn vị Độ (Degrees), nằm trong khoảng [0, 180].
    """
    a = np.array(a[:2])
    b = np.array(b[:2])
    c = np.array(c[:2])
    
    ba = a - b
    bc = c - b
    
    norm_ba = np.linalg.norm(ba)
    norm_bc = np.linalg.norm(bc)
    
    if norm_ba == 0 or norm_bc == 0:
        return 0.0
        
    cosine_angle = np.dot(ba, bc) / (norm_ba * norm_bc)
    angle = np.arccos(np.clip(cosine_angle, -1.0, 1.0))
    
    return np.degrees(angle)

def get_hand_state(hand_lms):
    """
    Calculates the 'state' of the hand (Open vs Closed) based on finger extension.
    Returns a score 0-1 (0: fully closed fist, 1: fully open palm).
    """
    if not hand_lms:
        return 0.5
    
    # Simple heuristic: distance between wrist (0) and tip of middle finger (12)
    # relative to the palm size (distance between 0 and 9)
    p0 = np.array(hand_lms[0][:2])
    p9 = np.array(hand_lms[9][:2])
    p12 = np.array(hand_lms[12][:2])
    
    palm_size = np.linalg.norm(p9 - p0) + 1e-6
    extension = np.linalg.norm(p12 - p0)
    
    return min(1.0, extension / (palm_size * 2))

def get_biomechanics_features(landmarks):
    """
    Trích xuất vector đặc trưng (Feature Vector) cho một frame hình.
    Bao gồm: 8 góc khớp xương chính (độ) và 2 trạng thái bàn tay (0-100).
    Scale hands lên 100 để có cùng độ lớn với đơn vị Độ.
    """
    if not landmarks or landmarks.get("pose") is None:
        return []
    
    pose = landmarks["pose"]
    
    # Key angles for Taekwondo
    angles = [
        calculate_angle(pose[11], pose[13], pose[15]), # L Elbow
        calculate_angle(pose[12], pose[14], pose[16]), # R Elbow
        calculate_angle(pose[23], pose[25], pose[27]), # L Knee
        calculate_angle(pose[24], pose[26], pose[28]), # R Knee
        calculate_angle(pose[13], pose[11], pose[23]), # L Shoulder
        calculate_angle(pose[14], pose[12], pose[24]), # R Shoulder
        calculate_angle(pose[11], pose[23], pose[25]), # L Hip
        calculate_angle(pose[12], pose[24], pose[26]), # R Hip
    ]
    
    # Hand states
    l_hand = get_hand_state(landmarks.get("left_hand"))
    r_hand = get_hand_state(landmarks.get("right_hand"))
    
    return angles + [l_hand * 100, r_hand * 100] # Scale hands to match angle magnitudes

class BiomechanicsMatcher:
    def __init__(self, reference_landmarks):
        self.ref_data = reference_landmarks
        self.ref_feature_seq = [get_biomechanics_features(f["landmarks"]) for f in reference_landmarks]
        self.ref_feature_seq = [f for f in self.ref_feature_seq if len(f) > 0]

    def find_start_frame(self, user_feature_seq, threshold=START_DETECTION_THRESHOLD):
        """
        Quét video của sinh viên để tìm frame đầu tiên khớp với thế đứng của video mẫu.
        Giúp loại bỏ khoảng thời gian 'chờ' không cần thiết ở đầu video.
        """
        if not user_feature_seq or not self.ref_feature_seq:
            return 0
            
        ref_start = np.array(self.ref_feature_seq[0])
        
        for i, u_feat in enumerate(user_feature_seq):
            dist = euclidean(u_feat, ref_start)
            if dist < threshold:
                return i
        return 0

    def compare(self, user_landmarks_seq):
        """
        Sử dụng thuật toán DTW (Dynamic Time Warping) để đồng bộ hóa và so sanh 2 chuỗi động tác.
        DTW cho phép so sánh chính xác ngay cả khi sinh viên thực hiện nhanh hay chậm hơn mẫu.
        Trả về: path (bản đồ khớp frame) và errors (chi tiết sai lệch từng cặp frame).
        """
        user_feature_seq = [get_biomechanics_features(f["landmarks"]) for f in user_landmarks_seq]
        user_feature_seq = [f for f in user_feature_seq if len(f) > 0]
        
        if not user_feature_seq or not self.ref_feature_seq:
            return None, None
            
        # Optimization: Find start frame to avoid comparing 'waiting' time
        start_idx = self.find_start_frame(user_feature_seq)
        trimmed_user_seq = user_feature_seq[start_idx:]
        
        distance, path = fastdtw(trimmed_user_seq, self.ref_feature_seq, dist=euclidean)
        
        # Adjust path indices back to original user sequence
        adjusted_path = [(u_idx + start_idx, r_idx) for u_idx, r_idx in path]

        # Calculate per-frame errors based on the path
        errors = []
        for u_idx_adj, r_idx in adjusted_path:
            # Need original untrimmed sequence for proper indexing if needed, 
            # but trimmed_user_seq corresponds to u_idx
            u_feat = np.array(user_feature_seq[u_idx_adj])
            r_feat = np.array(self.ref_feature_seq[r_idx])
            diff = np.abs(u_feat - r_feat)
            
            frame_error = {
                "total": np.mean(diff),
                "angles": diff[:8],
                "hands": diff[8:],
                "ref_idx": r_idx,
                "user_idx": u_idx_adj
            }
            errors.append(frame_error)

        return adjusted_path, errors

    def get_summary_score(self, errors):
        """
        Tính điểm tổng kết dựa trên sai số trung bình.
        Công thức: 100 - (sai số trung bình * 1.5). 
        Sai số trung bình 20 độ ứng với mức 50 điểm.
        """
        if not errors: return 0
        # Sai số trung bình 40 độ vẫn có thể đạt mức ~40 điểm (thay vì 0 như trước)
        avg_error = np.mean([e["total"] for e in errors])
        score = max(0, 100 - (avg_error * 1.5))
        return int(score)

def normalize_pose(pose_lms):
    """
    Translates and scales the pose so that the midpoint of hips is at (0,0)
    and the shoulder-to-hip distance is normalized to 1.0.
    """
    if not pose_lms: return None
    
    pose = np.array(pose_lms)
    # Midpoint of hips (23, 24)
    mid_hip = (pose[23][:2] + pose[24][:2]) / 2
    
    # Translate
    pose[:, :2] -= mid_hip
    
    # Scale based on torso height
    torso_height = np.linalg.norm(pose[11][:2] - pose[23][:2]) + 1e-6
    pose[:, :2] /= torso_height
    
    return pose.tolist()