"""
人体姿态估计与动作捕捉系统
支持实时视频处理和骨骼绘制
"""

import cv2
import numpy as np
from mmpose.apis import MMPoseInferencer
from typing import List, Tuple

class PoseCapture:
    def __init__(self, 
                 pose_model: str = 'rtmpose-m',
                 det_model: str = 'rtmdet_tiny',
                 device: str = 'cuda:0',
                 conf_threshold: float = 0.5):
        """
        初始化姿态捕捉器
        
        Args:
            pose_model: 姿态估计模型
            det_model: 人体检测模型
            device: 计算设备
            conf_threshold: 置信度阈值
        """
        self.inferencer = MMPoseInferencer(
            pose2d=pose_model,
            det_model=det_model,
            device=device
        )
        self.conf_threshold = conf_threshold
        
        # 骨骼连接关系（COCO 格式）
        self.skeleton = [
            (0, 1), (0, 2), (1, 3), (2, 4),    # 头部
            (5, 6),                             # 肩膀
            (5, 7), (7, 9),                    # 左臂
            (6, 8), (8, 10),                   # 右臂
            (5, 11), (6, 12),                  # 躯干
            (11, 13), (13, 15),                # 左腿
            (12, 14), (14, 16)                 # 右腿
        ]
        
        # 颜色映射
        self.colors = [
            (255, 0, 0), (255, 85, 0), (255, 170, 0),
            (255, 255, 0), (170, 255, 0), (85, 255, 0),
            (0, 255, 0), (0, 255, 85), (0, 255, 170),
            (0, 255, 255), (0, 170, 255), (0, 85, 255),
            (0, 0, 255), (85, 0, 255), (170, 0, 255),
            (255, 0, 255), (255, 0, 170)
        ]
    
    def process_frame(self, frame: np.ndarray) -> Tuple[np.ndarray, dict]:
        """
        处理单帧图像
        
        Args:
            frame: 输入图像
            
        Returns:
            绘制骨骼的图像，姿态数据
        """
        # 推理
        result = self.inferencer(frame)
        
        # 提取关键点
        if 'predictions' in result and len(result['predictions']) > 0:
            keypoints = result['predictions'][0]['keypoints']
            scores = result['predictions'][0]['keypoint_scores']
        else:
            keypoints = []
            scores = []
        
        # 绘制骨骼
        output_frame = self.draw_skeleton(frame.copy(), keypoints, scores)
        
        return output_frame, {
            'keypoints': keypoints,
            'scores': scores
        }
    
    def draw_skeleton(self, 
                      frame: np.ndarray, 
                      keypoints: List, 
                      scores: List) -> np.ndarray:
        """
        绘制人体骨骼
        
        Args:
            frame: 输入图像
            keypoints: 关键点列表
            scores: 置信度列表
            
        Returns:
            绘制骨骼后的图像
        """
        for person_idx, (kp, score) in enumerate(zip(keypoints, scores)):
            # 过滤低置信度关键点
            valid_mask = score > self.conf_threshold
            valid_kp = kp[valid_mask]
            
            # 绘制关键点
            for idx, point in enumerate(kp):
                if score[idx] > self.conf_threshold:
                    x, y = int(point[0]), int(point[1])
                    color = self.colors[idx % len(self.colors)]
                    cv2.circle(frame, (x, y), 5, color, -1)
            
            # 绘制骨骼连接
            for bone in self.skeleton:
                idx1, idx2 = bone
                if (idx1 < len(kp) and idx2 < len(kp) and
                    score[idx1] > self.conf_threshold and
                    score[idx2] > self.conf_threshold):
                    
                    pt1 = (int(kp[idx1][0]), int(kp[idx1][1]))
                    pt2 = (int(kp[idx2][0]), int(kp[idx2][1]))
                    color = self.colors[idx1 % len(self.colors)]
                    cv2.line(frame, pt1, pt2, color, 2)
        
        return frame
    
    def process_video(self, 
                      input_path: str, 
                      output_path: str,
                      show: bool = False):
        """
        处理视频文件
        
        Args:
            input_path: 输入视频路径
            output_path: 输出视频路径
            show: 是否实时显示
        """
        cap = cv2.VideoCapture(input_path)
        
        # 获取视频信息
        fps = cap.get(cv2.CAP_PROP_FPS)
        width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
        
        # 创建写入器
        fourcc = cv2.VideoWriter_fourcc(*'mp4v')
        out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
        
        frame_count = 0
        while True:
            ret, frame = cap.read()
            if not ret:
                break
            
            # 处理帧
            output_frame, pose_data = self.process_frame(frame)
            
            # 显示 FPS
            cv2.putText(output_frame, f'Frame: {frame_count}', 
                       (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
            
            # 写入输出
            out.write(output_frame)
            
            # 显示
            if show:
                cv2.imshow('Pose Capture', output_frame)
                if cv2.waitKey(1) & 0xFF == ord('q'):
                    break
            
            frame_count += 1
            if frame_count % 30 == 0:
                print(f"Processed {frame_count} frames")
        
        cap.release()
        out.release()
        cv2.destroyAllWindows()
        print(f"Video processed: {output_path}")


# 使用示例
if __name__ == '__main__':
    # 初始化
    capture = PoseCapture()
    
    # 处理单张图片
    import cv2
    image = cv2.imread('input.jpg')
    output, data = capture.process_frame(image)
    cv2.imwrite('output.jpg', output)
    
    # 处理视频
    # capture.process_video('input.mp4', 'output.mp4', show=True)
    
    # 处理摄像头
    # capture.process_video(0, 'webcam_output.mp4', show=True)