import cv2
import mediapipe as mp
import numpy as np
import time
import csv

# 初始化MediaPipe Pose模块
mp_pose = mp.solutions.pose
mp_drawing = mp.solutions.drawing_utils
mp_drawing_styles = mp.solutions.drawing_styles

# 存储角度数据
left_shoulder_angles = []  # 左肩关节角度
right_shoulder_angles = []  # 右肩关节角度
left_elbow_angles = []  # 左肘关节角度
right_elbow_angles = []  # 右肘关节角度
left_hip_angles = []  # 左髋关节角度
right_hip_angles = []  # 右髋关节角度
left_knee_angles = []  # 左膝关节角度
right_knee_angles = []  # 右膝关节角度
left_ankle_angles = []  # 左踝关节角度
right_ankle_angles = []  # 右踝关节角度
left_wrist_angles = []  # 左腕关节角度
right_wrist_angles = []  # 右腕关节角度

# 存储关键点数据
keypoints_data = []  # 每一帧的关键点数据

def calculate_angle(a, b, c):
    """
    计算三个点之间的夹角（以度为单位）。
    a, b, c 是三个点的坐标，b 是顶点。
    """
    a = np.array(a)
    b = np.array(b)
    c = np.array(c)

    # 计算向量
    ba = a - b
    bc = c - b

    # 计算夹角
    cosine_angle = np.dot(ba, bc) / (np.linalg.norm(ba) * np.linalg.norm(bc))
    angle = np.arccos(cosine_angle)
    return np.degrees(angle)

# 使用本地视频推理
cap = cv2.VideoCapture(r"G:\运动方案\动作评价\卧推.mp4")
with mp_pose.Pose(
        min_detection_confidence=0.5,
        min_tracking_confidence=0.5,
        model_complexity=1) as pose:

    frame_count = 0
    while cap.isOpened():
        success, image = cap.read()
        if not success:
            print("Video ended or corrupted.")
            break  # 视频结束或损坏时退出

        # 镜像翻转视频帧
        image = cv2.flip(image, 1)  # 水平翻转图像

        start = time.time()

        # 提高性能
        image.flags.writeable = False
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        results = pose.process(image)

        # 绘制姿态注释
        image.flags.writeable = True
        image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)

        mp_drawing.draw_landmarks(
            image,
            results.pose_landmarks,
            mp_pose.POSE_CONNECTIONS,
            landmark_drawing_spec=mp_drawing_styles.get_default_pose_landmarks_style())

        end = time.time()
        fps = 1 / (end - start)
        fps_text = f"FPS: {fps:.2f}"
        # 实时显示帧数
        cv2.putText(image, fps_text, (100, 50),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.75, (255, 255, 255), 3)

        # 如果检测到世界坐标中的关键点，计算关节角度
        if results.pose_world_landmarks:
            # 获取关键点数据
            frame_keypoints = []
            for landmark in results.pose_world_landmarks.landmark:
                frame_keypoints.append([landmark.x, landmark.y, landmark.z])
            keypoints_data.append(frame_keypoints)

            # 计算左肩关节角度
            left_shoulder = [results.pose_world_landmarks.landmark[mp_pose.PoseLandmark.LEFT_SHOULDER].x,
                             results.pose_world_landmarks.landmark[mp_pose.PoseLandmark.LEFT_SHOULDER].y,
                             results.pose_world_landmarks.landmark[mp_pose.PoseLandmark.LEFT_SHOULDER].z]
            left_elbow = [results.pose_world_landmarks.landmark[mp_pose.PoseLandmark.LEFT_ELBOW].x,
                          results.pose_world_landmarks.landmark[mp_pose.PoseLandmark.LEFT_ELBOW].y,
                          results.pose_world_landmarks.landmark[mp_pose.PoseLandmark.LEFT_ELBOW].z]
            left_hip = [results.pose_world_landmarks.landmark[mp_pose.PoseLandmark.LEFT_HIP].x,
                        results.pose_world_landmarks.landmark[mp_pose.PoseLandmark.LEFT_HIP].y,
                        results.pose_world_landmarks.landmark[mp_pose.PoseLandmark.LEFT_HIP].z]

            left_shoulder_angle = calculate_angle(left_elbow, left_shoulder, left_hip)
            left_shoulder_angles.append(left_shoulder_angle)

            # 计算右肩关节角度
            right_shoulder = [results.pose_world_landmarks.landmark[mp_pose.PoseLandmark.RIGHT_SHOULDER].x,
                              results.pose_world_landmarks.landmark[mp_pose.PoseLandmark.RIGHT_SHOULDER].y,
                              results.pose_world_landmarks.landmark[mp_pose.PoseLandmark.RIGHT_SHOULDER].z]
            right_elbow = [results.pose_world_landmarks.landmark[mp_pose.PoseLandmark.RIGHT_ELBOW].x,
                           results.pose_world_landmarks.landmark[mp_pose.PoseLandmark.RIGHT_ELBOW].y,
                           results.pose_world_landmarks.landmark[mp_pose.PoseLandmark.RIGHT_ELBOW].z]
            right_hip = [results.pose_world_landmarks.landmark[mp_pose.PoseLandmark.RIGHT_HIP].x,
                         results.pose_world_landmarks.landmark[mp_pose.PoseLandmark.RIGHT_HIP].y,
                         results.pose_world_landmarks.landmark[mp_pose.PoseLandmark.RIGHT_HIP].z]

            right_shoulder_angle = calculate_angle(right_elbow, right_shoulder, right_hip)
            right_shoulder_angles.append(right_shoulder_angle)

            # 计算左肘关节角度
            left_wrist = [results.pose_world_landmarks.landmark[mp_pose.PoseLandmark.LEFT_WRIST].x,
                          results.pose_world_landmarks.landmark[mp_pose.PoseLandmark.LEFT_WRIST].y,
                          results.pose_world_landmarks.landmark[mp_pose.PoseLandmark.LEFT_WRIST].z]

            left_elbow_angle = calculate_angle(left_shoulder, left_elbow, left_wrist)
            left_elbow_angles.append(left_elbow_angle)

            # 计算右肘关节角度
            right_wrist = [results.pose_world_landmarks.landmark[mp_pose.PoseLandmark.RIGHT_WRIST].x,
                           results.pose_world_landmarks.landmark[mp_pose.PoseLandmark.RIGHT_WRIST].y,
                           results.pose_world_landmarks.landmark[mp_pose.PoseLandmark.RIGHT_WRIST].z]

            right_elbow_angle = calculate_angle(right_shoulder, right_elbow, right_wrist)
            right_elbow_angles.append(right_elbow_angle)

            # 计算左髋关节角度
            left_knee = [results.pose_world_landmarks.landmark[mp_pose.PoseLandmark.LEFT_KNEE].x,
                         results.pose_world_landmarks.landmark[mp_pose.PoseLandmark.LEFT_KNEE].y,
                         results.pose_world_landmarks.landmark[mp_pose.PoseLandmark.LEFT_KNEE].z]

            left_hip_angle = calculate_angle(left_shoulder, left_hip, left_knee)
            left_hip_angles.append(left_hip_angle)

            # 计算右髋关节角度
            right_knee = [results.pose_world_landmarks.landmark[mp_pose.PoseLandmark.RIGHT_KNEE].x,
                          results.pose_world_landmarks.landmark[mp_pose.PoseLandmark.RIGHT_KNEE].y,
                          results.pose_world_landmarks.landmark[mp_pose.PoseLandmark.RIGHT_KNEE].z]

            right_hip_angle = calculate_angle(right_shoulder, right_hip, right_knee)
            right_hip_angles.append(right_hip_angle)

            # 计算左膝关节角度
            left_ankle = [results.pose_world_landmarks.landmark[mp_pose.PoseLandmark.LEFT_ANKLE].x,
                          results.pose_world_landmarks.landmark[mp_pose.PoseLandmark.LEFT_ANKLE].y,
                          results.pose_world_landmarks.landmark[mp_pose.PoseLandmark.LEFT_ANKLE].z]

            left_knee_angle = calculate_angle(left_hip, left_knee, left_ankle)
            left_knee_angles.append(left_knee_angle)

            # 计算右膝关节角度
            right_ankle = [results.pose_world_landmarks.landmark[mp_pose.PoseLandmark.RIGHT_ANKLE].x,
                           results.pose_world_landmarks.landmark[mp_pose.PoseLandmark.RIGHT_ANKLE].y,
                           results.pose_world_landmarks.landmark[mp_pose.PoseLandmark.RIGHT_ANKLE].z]

            right_knee_angle = calculate_angle(right_hip, right_knee, right_ankle)
            right_knee_angles.append(right_knee_angle)

            # 计算左踝关节角度
            left_foot_index = [results.pose_world_landmarks.landmark[mp_pose.PoseLandmark.LEFT_FOOT_INDEX].x,
                               results.pose_world_landmarks.landmark[mp_pose.PoseLandmark.LEFT_FOOT_INDEX].y,
                               results.pose_world_landmarks.landmark[mp_pose.PoseLandmark.LEFT_FOOT_INDEX].z]

            left_ankle_angle = calculate_angle(left_knee, left_ankle, left_foot_index)
            left_ankle_angles.append(left_ankle_angle)

            # 计算右踝关节角度
            right_foot_index = [results.pose_world_landmarks.landmark[mp_pose.PoseLandmark.RIGHT_FOOT_INDEX].x,
                                results.pose_world_landmarks.landmark[mp_pose.PoseLandmark.RIGHT_FOOT_INDEX].y,
                                results.pose_world_landmarks.landmark[mp_pose.PoseLandmark.RIGHT_FOOT_INDEX].z]

            right_ankle_angle = calculate_angle(right_knee, right_ankle, right_foot_index)
            right_ankle_angles.append(right_ankle_angle)

            # 计算左腕关节角度
            left_wrist_angle = calculate_angle(left_elbow, left_wrist, left_shoulder)
            left_wrist_angles.append(left_wrist_angle)

            # 计算右腕关节角度
            right_wrist_angle = calculate_angle(right_elbow, right_wrist, right_shoulder)
            right_wrist_angles.append(right_wrist_angle)

            # 输出角度数据
            print(f"Frame {frame_count}:")
            print(f"Left Shoulder Angle: {left_shoulder_angle:.2f} degrees")
            print(f"Right Shoulder Angle: {right_shoulder_angle:.2f} degrees")
            print(f"Left Elbow Angle: {left_elbow_angle:.2f} degrees")
            print(f"Right Elbow Angle: {right_elbow_angle:.2f} degrees")
            print(f"Left Hip Angle: {left_hip_angle:.2f} degrees")
            print(f"Right Hip Angle: {right_hip_angle:.2f} degrees")
            print(f"Left Knee Angle: {left_knee_angle:.2f} degrees")
            print(f"Right Knee Angle: {right_knee_angle:.2f} degrees")
            print(f"Left Ankle Angle: {left_ankle_angle:.2f} degrees")
            print(f"Right Ankle Angle: {right_ankle_angle:.2f} degrees")
            print(f"Left Wrist Angle: {left_wrist_angle:.2f} degrees")
            print(f"Right Wrist Angle: {right_wrist_angle:.2f} degrees")
            print("-" * 40)

        # 显示视频帧
        cv2.imshow('MediaPipe Pose', image)
        if cv2.waitKey(5) & 0xFF == 27:
            break

        frame_count += 1

cap.release()
cv2.destroyAllWindows()

# 保存数据为 .npz 文件
# output_npz_path = r"C:\Users\messi10\Desktop\运动方案\动作评价\八段锦\wotui.npz"
# np.savez(
#     output_npz_path,
#     left_shoulder_angles=left_shoulder_angles,
#     right_shoulder_angles=right_shoulder_angles,
#     left_elbow_angles=left_elbow_angles,
#     right_elbow_angles=right_elbow_angles,
#     left_hip_angles=left_hip_angles,
#     right_hip_angles=right_hip_angles,
#     left_knee_angles=left_knee_angles,
#     right_knee_angles=right_knee_angles,
#     left_ankle_angles=left_ankle_angles,
#     right_ankle_angles=right_ankle_angles,
#     left_wrist_angles=left_wrist_angles,
#     right_wrist_angles=right_wrist_angles,
#     keypoints_data=keypoints_data
# )
# print(f"Joint data saved to: {output_npz_path}")

# 保存数据为 .csv 文件
output_csv_path = r"F:\pythonzfy\动作评价\velocity\wotui.csv"
with open(output_csv_path, mode='w', newline='') as file:
    writer = csv.writer(file)
    # 写入表头
    writer.writerow([
        "Frame",
        "Left Shoulder Angle", "Right Shoulder Angle",
        "Left Elbow Angle", "Right Elbow Angle",
        "Left Hip Angle", "Right Hip Angle",
        "Left Knee Angle", "Right Knee Angle",
        "Left Ankle Angle", "Right Ankle Angle",
        "Left Wrist Angle", "Right Wrist Angle",
        "Keypoints"
    ])
    # 写入每一帧的数据
    for i in range(len(left_shoulder_angles)):
        writer.writerow([
            i,
            left_shoulder_angles[i], right_shoulder_angles[i],
            left_elbow_angles[i], right_elbow_angles[i],
            left_hip_angles[i], right_hip_angles[i],
            left_knee_angles[i], right_knee_angles[i],
            left_ankle_angles[i], right_ankle_angles[i],
            left_wrist_angles[i], right_wrist_angles[i],
            keypoints_data[i]
        ])
print(f"Joint data saved to: {output_csv_path}")改成只要肩肘腕关节的数据