import cv2
import numpy as np

# Mapping from feature vector indices (from compare.py) to MediaPipe landmark IDs
FEATURE_TO_LANDMARK = {
    0: 13, # L Elbow
    1: 14, # R Elbow
    2: 25, # L Knee
    3: 26, # R Knee
    4: 11, # L Shoulder
    5: 12, # R Shoulder
    6: 23, # L Hip
    7: 24, # R Hip
    8: 15, # L Hand (Wrist)
    9: 16  # R Hand (Wrist)
}

FEATURE_LABELS = {
    0: "Tay trái", 1: "Tay phải",
    2: "Chân trái", 3: "Chân phải",
    4: "Vai trái", 5: "Vai phải",
    6: "Hông trái", 7: "Hông phải",
    8: "Bàn tay trái", 9: "Bàn tay phải"
}

def draw_error_highlight(frame, landmark, color=(0, 0, 255), radius=35, thickness=4):
    """
    Draws a red circle around a faulty landmark.
    """
    h, w, _ = frame.shape
    cx, cy = int(landmark[0] * w), int(landmark[1] * h)
    cv2.circle(frame, (cx, cy), radius, color, thickness)
    cv2.circle(frame, (cx, cy), radius + 5, color, 1)

def overlay_skeleton(frame, landmarks, color=(0, 255, 0), thickness=2):
    """
    Draws skeleton connections on the frame.
    """
    h, w, _ = frame.shape
    pose = landmarks.get("pose")
    if not pose: return frame
    
    connections = [
        (11, 13), (13, 15), (12, 14), (14, 16),
        (11, 12), (23, 24), (11, 23), (12, 24),
        (23, 25), (25, 27), (24, 26), (26, 28)
    ]
    
    for s, e in connections:
        if s < len(pose) and e < len(pose):
            p1 = (int(pose[s][0] * w), int(pose[s][1] * h))
            p2 = (int(pose[e][0] * w), int(pose[e][1] * h))
            cv2.line(frame, p1, p2, color, thickness)
    return frame

def create_combined_overlay(user_frame, user_landmarks, frame_error=None, error_threshold=120.0):
    """
    Vẽ các chỉ số sửa lỗi trực tiếp lên frame của sinh viên.
    error_threshold=120.0: Ngưỡng cộng dồn sai lệch góc để hiển thị vòng tròn đỏ.
    """
    output_frame = user_frame.copy()
    
    # Vẽ khung xương của sinh viên (màu xanh lá - color=(0, 255, 0))
    output_frame = overlay_skeleton(output_frame, user_landmarks, color=(0, 255, 0), thickness=2)
    
    # Kiểm tra và làm nổi bật các vị trí lỗi dựa trên kết quả so sánh
    if frame_error and user_landmarks.get("pose"):
        pose = user_landmarks["pose"]
        
        # Duyệt qua các sai lệch góc khớp
        for i, diff in enumerate(frame_error.get("angles", [])):
            if diff > error_threshold:
                if i in FEATURE_TO_LANDMARK:
                    draw_error_highlight(output_frame, pose[FEATURE_TO_LANDMARK[i]])
                    
        # Duyệt qua sai lệch trạng thái bàn tay (ngưỡng fix cứng 30.0)
        for i, diff in enumerate(frame_error.get("hands", [])):
            if diff > 30.0:
                if i+8 in FEATURE_TO_LANDMARK:
                    # Màu cam (255, 128, 0) dành riêng cho cảnh báo bàn tay
                    draw_error_highlight(output_frame, pose[FEATURE_TO_LANDMARK[i+8]], color=(255, 128, 0))
    
    return output_frame

def generate_result_video(user_video_path, user_landmarks_seq, path_mapping, errors, output_path):
    """
    Renders the study video showing ONLY the student but with AI highlights.
    """
    cap_user = cv2.VideoCapture(user_video_path)
    fps = cap_user.get(cv2.CAP_PROP_FPS) or 30.0
    w = int(cap_user.get(cv2.CAP_PROP_FRAME_WIDTH))
    h = int(cap_user.get(cv2.CAP_PROP_FRAME_HEIGHT))
    
    fourcc = cv2.VideoWriter_fourcc(*'vp80') 
    out = cv2.VideoWriter(output_path, fourcc, fps, (w, h))
    
    curr_u = 0
    while cap_user.isOpened():
        ret, u_frame = cap_user.read()
        if not ret: break
        
        # Find matches in path to get errors for this frame
        matches = [p_idx for p_idx, (u, r) in enumerate(path_mapping) if u == curr_u]
        
        if matches:
            match_idx = matches[0]
            user_lms = user_landmarks_seq[curr_u]["landmarks"]
            frame_err = errors[match_idx]
            
            # Create overlay on user frame only
            combined = create_combined_overlay(u_frame, user_lms, frame_err)
            out.write(combined)
        else:
            # If no match (e.g. before start frame), just write raw frame or skip
            out.write(u_frame)
            
        curr_u += 1
        
    cap_user.release()
    out.release()