# 核心依赖
import cv2
import numpy as np
import mediapipe as mp
import matplotlib.pyplot as plt
from typing import List, Tuple, Dict
import pandas as pd
from pathlib import Path

# OpenCV
pip install opencv-python opencv-contrib-python

# MediaPipe
pip install mediapipe

# 其他依赖
pip install numpy matplotlib pandas

# 初始化MediaPipe姿态检测
mp_pose = mp.solutions.pose
mp_drawing = mp.solutions.drawing_utils
mp_drawing_styles = mp.solutions.drawing_styles

# 创建姿态检测器
pose = mp_pose.Pose(
    static_image_mode=False,      # 视频流模式
    model_complexity=2,           # 模型复杂度 (0, 1, 2)
    smooth_landmarks=True,        # 平滑关键点
    min_detection_confidence=0.5, # 检测置信度阈值
    min_tracking_confidence=0.5   # 追踪置信度阈值
)

class RunningVideoProcessor:
    """跑步视频处理器"""

    def __init__(self, video_path: str):
        self.video_path = video_path
        self.cap = cv2.VideoCapture(video_path)
        self.fps = self.cap.get(cv2.CAP_PROP_FPS)
        self.frame_count = int(self.cap.get(cv2.CAP_PROP_FRAME_COUNT))
        self.width = int(self.cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        self.height = int(self.cap.get(cv2.CAP_PROP_FRAME_HEIGHT))

        print(f"视频信息:")
        print(f"  帧率: {self.fps} FPS")
        print(f"  总帧数: {self.frame_count}")
        print(f"  分辨率: {self.width}x{self.height}")

    def read_frame(self, frame_number: int):
        """读取指定帧"""
        self.cap.set(cv2.CAP_PROP_POS_FRAMES, frame_number)
        ret, frame = self.cap.read()
        if not ret:
            raise ValueError(f"无法读取帧 {frame_number}")
        return frame

    def extract_frames(self, interval: int = 5):
        """
        提取帧序列

        参数:
            interval: 提取间隔（每N帧提取1帧）
        """
        frames = []
        frame_indices = []

        for idx in range(0, self.frame_count, interval):
            frame = self.read_frame(idx)
            frames.append(frame)
            frame_indices.append(idx)

        return frames, frame_indices

    def release(self):
        """释放资源"""
        self.cap.release()

# 使用示例
processor = RunningVideoProcessor("running_video.mp4")
frames, indices = processor.extract_frames(interval=10)  # 每10帧提取1帧

def detect_pose_landmarks(frame, pose_detector):
    """
    检测单帧中的姿态关键点

    MediaPipe检测33个3D关键点:
    - 0: 鼻子
    - 11-12: 左右肩
    - 23-24: 左右髋
    - 25-26: 左右膝
    - 27-28: 左右踝
    - 29-30: 左右脚跟
    - 31-32: 左右脚尖
    """
    # 转换为RGB
    rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)

    # 检测姿态
    results = pose_detector.process(rgb_frame)

    return results

# 测试关键点检测
frame = processor.read_frame(0)
results = detect_pose_landmarks(frame, pose)

if results.pose_landmarks:
    print("检测到姿态关键点!")
    # 关键点数据
    landmarks = results.pose_landmarks.landmark
    print(f"关键点数量: {len(landmarks)}")

def extract_body_landmarks(landmarks) -> Dict[str, Tuple[float, float]]:
    """
    提取关键身体部位的坐标

    返回字典:
        {
            'nose': (x, y),
            'left_shoulder': (x, y),
            'right_shoulder': (x, y),
            'left_hip': (x, y),
            'right_hip': (x, y),
            'left_knee': (x, y),
            'right_knee': (x, y),
            'left_ankle': (x, y),
            'right_ankle': (x, y),
        }
    """
    key_points = {
        'nose': landmarks[0],
        'left_shoulder': landmarks[11],
        'right_shoulder': landmarks[12],
        'left_hip': landmarks[23],
        'right_hip': landmarks[24],
        'left_knee': landmarks[25],
        'right_knee': landmarks[26],
        'left_ankle': landmarks[27],
        'right_ankle': landmarks[28],
        'left_heel': landmarks[29],
        'right_heel': landmarks[30],
        'left_foot_index': landmarks[31],
        'right_foot_index': landmarks[32],
    }

    # 转换为像素坐标（归一化坐标 × 图像尺寸）
    body_parts = {}
    for name, landmark in key_points.items():
        x = landmark.x * processor.width
        y = landmark.y * processor.height
        body_parts[name] = (x, y)

    return body_parts

def calculate_angle(a: Tuple[float, float],
                   b: Tuple[float, float],
                   c: Tuple[float, float]) -> float:
    """
    计算三点形成的角度（ABC，B为顶点）

    使用余弦定理:
    cos(θ) = (AB² + BC² - AC²) / (2 × AB × BC)
    """
    # 转换为numpy数组
    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(np.clip(cosine_angle, -1.0, 1.0))

    # 转换为度数
    angle_degrees = np.degrees(angle)

    return angle_degrees

def analyze_joint_angles(body_parts: Dict[str, Tuple[float, float]]) -> Dict[str, float]:
    """分析关键关节角度"""

    angles = {}

    # 1. 左膝角度（髋-膝-踝）
    if all(k in body_parts for k in ['left_hip', 'left_knee', 'left_ankle']):
        angles['left_knee'] = calculate_angle(
            body_parts['left_hip'],
            body_parts['left_knee'],
            body_parts['left_ankle']
        )

    # 2. 右膝角度
    if all(k in body_parts for k in ['right_hip', 'right_knee', 'right_ankle']):
        angles['right_knee'] = calculate_angle(
            body_parts['right_hip'],
            body_parts['right_knee'],
            body_parts['right_ankle']
        )

    # 3. 左髋角度（肩-髋-膝）
    if all(k in body_parts for k in ['left_shoulder', 'left_hip', 'left_knee']):
        angles['left_hip'] = calculate_angle(
            body_parts['left_shoulder'],
            body_parts['left_hip'],
            body_parts['left_knee']
        )

    # 4. 右髋角度
    if all(k in body_parts for k in ['right_shoulder', 'right_hip', 'right_knee']):
        angles['right_hip'] = calculate_angle(
            body_parts['right_shoulder'],
            body_parts['right_hip'],
            body_parts['right_knee']
        )

    # 5. 躯干倾斜角（肩-髋-垂直）
    if all(k in body_parts for k in ['left_shoulder', 'right_shoulder', 'left_hip']):
        # 计算肩线中点
        shoulder_mid = (
            (body_parts['left_shoulder'][0] + body_parts['right_shoulder'][0]) / 2,
            (body_parts['left_shoulder'][1] + body_parts['right_shoulder'][1]) / 2
        )
        hip_point = body_parts['left_hip']

        # 与垂直线的夹角
        vertical = (shoulder_mid[0], shoulder_mid[1] - 100)
        angles['trunk_lean'] = calculate_angle(hip_point, shoulder_mid, vertical) - 90

    return angles

class RunningFormAnalyzer:
    """跑步姿态分析器"""

    def __init__(self):
        self.assessments = []

    def assess_form(self, angles: Dict[str, float]) -> Dict[str, any]:
        """
        评估跑步姿态

        检查项目:
        1. 膝关节伸展（着地时）
        2. 髋关节活动范围
        3. 躯干倾斜
        4. 左右对称性
        """
        assessment = {
            'overall_score': 100,
            'issues': [],
            'recommendations': []
        }

        # 1. 膝关节评估
        for side in ['left', 'right']:
            knee_key = f'{side}_knee'
            if knee_key in angles:
                knee_angle = angles[knee_key]

                # 着地时膝角应 > 140°（充分伸展）
                if knee_angle ＜ 140:
                    assessment['overall_score'] -= 10
                    assessment['issues'].append(
                        f"{side}膝关节伸展不足 ({knee_angle:.1f}°)"
                    )
                    assessment['recommendations'].append(
                        f"加强{side}腿股四头肌力量"
                    )

        # 2. 髋关节活动范围
        for side in ['left', 'right']:
            hip_key = f'{side}_hip'
            if hip_key in angles:
                hip_angle = angles[hip_key]

                # 理想范围: 120-150°
                if hip_angle ＜ 120:
                    assessment['overall_score'] -= 10
                    assessment['issues'].append(
                        f"{side}髋关节活动范围不足 ({hip_angle:.1f}°)"
                    )
                    assessment['recommendations'].append(
                        f"增加{side}髋灵活性训练"
                    )

        # 3. 躯干倾斜评估
        if 'trunk_lean' in angles:
            trunk_angle = angles['trunk_lean']

            # 理想: 轻微前倾 (5-10°)
            if trunk_angle ＜ 0:  # 后倾
                assessment['overall_score'] -= 15
                assessment['issues'].append(
                    f"躯干后倾 ({trunk_angle:.1f}°)"
                )
                assessment['recommendations'].append(
                    "练习核心稳定性，保持轻微前倾"
                )
            elif trunk_angle > 15:  # 过度前倾
                assessment['overall_score'] -= 10
                assessment['issues'].append(
                    f"躯干过度前倾 ({trunk_angle:.1f}°)"
                )
                assessment['recommendations'].append(
                    "加强背肌力量，控制前倾角度"
                )

        # 4. 左右对称性评估
        if 'left_knee' in angles and 'right_knee' in angles:
            knee_diff = abs(angles['left_knee'] - angles['right_knee'])
            if knee_diff > 15:
                assessment['overall_score'] -= 10
                assessment['issues'].append(
                    f"左右不对称: 膝角差异 {knee_diff:.1f}°"
                )
                assessment['recommendations'].append(
                    "进行单腿平衡和力量练习，改善不对称"
                )

        return assessment

    def analyze_video_sequence(self, processor, frames):
        """分析整个视频序列"""
        sequence_data = []

        for i, frame in enumerate(frames):
            # 检测姿态
            results = detect_pose_landmarks(frame, pose)

            if results.pose_landmarks:
                # 提取关键点
                body_parts = extract_body_landmarks(results.pose_landmarks.landmark)

                # 计算角度
                angles = analyze_joint_angles(body_parts)

                # 评估姿态
                assessment = self.assess_form(angles)

                frame_data = {
                    'frame': i,
                    'angles': angles,
                    'assessment': assessment,
                    'body_parts': body_parts
                }
                sequence_data.append(frame_data)

        return sequence_data

def draw_skeleton(frame, body_parts: Dict[str, Tuple[float, float]],
                 assessment: Dict[str, any]):
    """在帧上绘制骨架和评估结果"""

    frame_copy = frame.copy()

    # 定义骨架连接
    connections = [
        ('left_shoulder', 'right_shoulder'),
        ('left_shoulder', 'left_hip'),
        ('right_shoulder', 'right_hip'),
        ('left_hip', 'right_hip'),
        ('left_hip', 'left_knee'),
        ('left_knee', 'left_ankle'),
        ('right_hip', 'right_knee'),
        ('right_knee', 'right_ankle'),
    ]

    # 绘制连接线
    for start, end in connections:
        if start in body_parts and end in body_parts:
            cv2.line(
                frame_copy,
                (int(body_parts[start][0]), int(body_parts[start][1])),
                (int(body_parts[end][0]), int(body_parts[end][1])),
                (0, 255, 0), 2
            )

    # 绘制关键点
    for name, (x, y) in body_parts.items():
        cv2.circle(frame_copy, (int(x), int(y)), 5, (0, 0, 255), -1)

    # 显示评分
    score = assessment['overall_score']
    color = (0, 255, 0) if score >= 80 else (0, 165, 255) if score >= 60 else (0, 0, 255)
    cv2.putText(
        frame_copy,
        f"Score: {score}",
        (10, 30),
        cv2.FONT_HERSHEY_SIMPLEX,
        1, color, 2
    )

    # 显示主要问题
    y_offset = 70
    for issue in assessment['issues'][:3]:  # 最多显示3个问题
        cv2.putText(
            frame_copy,
            f"- {issue}",
            (10, y_offset),
            cv2.FONT_HERSHEY_SIMPLEX,
            0.5, (0, 0, 255), 1
        )
        y_offset += 25

    return frame_copy

def generate_analysis_report(sequence_data: List[Dict]) -> Dict:
    """生成分析报告"""

    # 计算平均角度
    angle_names = ['left_knee', 'right_knee', 'left_hip', 'right_hip', 'trunk_lean']
    angle_stats = {}

    for angle_name in angle_names:
        values = [
            frame['angles'].get(angle_name, None)
            for frame in sequence_data
            if angle_name in frame['angles']
        ]
        if values:
            angle_stats[angle_name] = {
                'mean': np.mean(values),
                'std': np.std(values),
                'min': np.min(values),
                'max': np.max(values)
            }

    # 计算平均得分
    scores = [frame['assessment']['overall_score'] for frame in sequence_data]

    # 统计问题频率
    issue_counts = {}
    for frame in sequence_data:
        for issue in frame['assessment']['issues']:
            issue_counts[issue] = issue_counts.get(issue, 0) + 1

    # 找出最常见的问题
    top_issues = sorted(issue_counts.items(), key=lambda x: x[1], reverse=True)[:5]

    report = {
        'overall_score': np.mean(scores),
        'score_range': (np.min(scores), np.max(scores)),
        'angle_statistics': angle_stats,
        'common_issues': top_issues,
        'total_frames_analyzed': len(sequence_data)
    }

    return report

# 生成报告
analyzer = RunningFormAnalyzer()
sequence_data = analyzer.analyze_video_sequence(processor, frames)
report = generate_analysis_report(sequence_data)

# 打印报告
print(f"=== 跑步姿态分析报告 ===")
print(f"总体评分: {report['overall_score']:.1f}/100")
print(f"评分范围: {report['score_range'][0]:.0f} - {report['score_range'][1]:.0f}")
print(f"\n常见问题:")
for issue, count in report['common_issues']:
    print(f"  - {issue} (出现 {count} 次)")

def create_analysis_video(input_path: str, output_path: str,
                         sequence_data: List[Dict],
                         processor: RunningVideoProcessor):
    """创建带标注的分析视频"""

    # 获取视频信息
    fps = processor.fps
    width = processor.width
    height = processor.height

    # 创建视频写入器
    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))

    for frame_data in sequence_data:
        # 读取原始帧
        frame = processor.read_frame(frame_data['frame'])

        # 绘制分析结果
        annotated_frame = draw_skeleton(
            frame,
            frame_data['body_parts'],
            frame_data['assessment']
        )

        # 写入输出视频
        out.write(annotated_frame)

    out.release()
    print(f"分析视频已保存: {output_path}")

# 使用示例
create_analysis_video(
    "running_video.mp4",
    "running_analysis.mp4",
    sequence_data,
    processor
)