import numpy as np
from dataclasses import dataclass
from typing import List, Tuple, Optional

"""
IMS OOP (Out-of-Position) 异常姿态检测
基于 VideoRun2D 无标记 2D 姿态追踪

参考: arXiv:2608.19480 VideoRun2D
"""

# COCO 17 关键点
KEYPOINT_NAMES = [
    'nose', 'left_eye', 'right_eye', 'left_ear', 'right_ear',
    'left_shoulder', 'right_shoulder', 'left_elbow', 'right_elbow',
    'left_wrist', 'right_wrist', 'left_hip', 'right_hip',
    'left_knee', 'right_knee', 'left_ankle', 'right_ankle'
]

@dataclass
class Keypoint:
    name: str
    x: float
    y: float
    confidence: float

@dataclass
class PoseFrame:
    """单帧姿态"""
    keypoints: List[Keypoint]
    timestamp: float
    
    def get(self, name: str) -> Optional[Keypoint]:
        for kp in self.keypoints:
            if kp.name == name:
                return kp
        return None
    
    def shoulder_midpoint(self) -> Tuple[float, float]:
        ls = self.get('left_shoulder')
        rs = self.get('right_shoulder')
        if ls and rs:
            return ((ls.x + rs.x) / 2, (ls.y + rs.y) / 2)
        return (0, 0)
    
    def hip_midpoint(self) -> Tuple[float, float]:
        lh = self.get('left_hip')
        rh = self.get('right_hip')
        if lh and rh:
            return ((lh.x + rh.x) / 2, (lh.y + rh.y) / 2)
        return (0, 0)


class OOPDetector:
    """OOP 异常姿态检测器"""
    
    # OOP 姿态定义
    OOP_TYPES = {
        'normal': '正常坐姿',
        'leaning_forward': '前倾（看手机/低头）',
        'leaning_side': '侧倾（靠窗口/睡觉）',
        'slouching': '瘫坐（脊柱弯曲）',
        'feet_dashboard': '脚搭仪表台',
        'turned_back': '背对前方（回头照顾后座）',
        'sleeping': '睡眠',
        'unreachable': '无法检测'
    }
    
    def __init__(self, image_w: int = 640, image_h: int = 480):
        self.w = image_w
        self.h = image_h
    
    def detect_posture(self, frame: PoseFrame) -> dict:
        """检测乘员姿态"""
        nose = frame.get('nose')
        ls = frame.get('left_shoulder')
        rs = frame.get('right_shoulder')
        lh = frame.get('left_hip')
        rh = frame.get('right_hip')
        lk = frame.get('left_knee')
        rk = frame.get('right_knee')
        la = frame.get('left_ankle')
        ra = frame.get('right_ankle')
        
        if not all([nose, ls, rs, lh, rh]):
            return {
                'posture': 'unreachable',
                'confidence': 0.0,
                'details': '关键点不完整'
            }
        
        # 躯干角度（肩-髋连线与垂直线的夹角）
        shoulder_mid = frame.shoulder_midpoint()
        hip_mid = frame.hip_midpoint()
        
        torso_dx = shoulder_mid[0] - hip_mid[0]
        torso_dy = shoulder_mid[1] - hip_mid[1]
        torso_angle = np.degrees(np.arctan2(torso_dx, -torso_dy))
        
        # 头部位置（相对于躯干）
        head_offset_x = nose.x - shoulder_mid[0]
        head_offset_y = nose.y - shoulder_mid[1]
        head_angle = np.degrees(np.arctan2(head_offset_x, -head_offset_y))
        
        # 肩膀水平度
        shoulder_tilt = np.degrees(np.arctan2(
            rs.y - ls.y, rs.x - ls.x))
        
        # 膝盖位置（是否在仪表台方向）
        knee_height_ratio = 0
        if lk and rk:
            knee_y = (lk.y + rk.y) / 2
            knee_height_ratio = 1 - (knee_y / self.h)
        
        # 判断
        confidence = 0.8
        details = {
            'torso_angle': round(torso_angle, 1),
            'head_angle': round(head_angle, 1),
            'shoulder_tilt': round(shoulder_tilt, 1),
            'knee_height_ratio': round(knee_height_ratio, 2)
        }
        
        if abs(torso_angle) < 10 and abs(shoulder_tilt) < 10:
            posture = 'normal'
        elif torso_angle < -15:
            posture = 'leaning_forward'
        elif abs(shoulder_tilt) > 15:
            posture = 'leaning_side'
        elif torso_angle > 20:
            posture = 'slouching'
        elif knee_height_ratio > 0.7:
            posture = 'feet_dashboard'
        elif abs(head_angle) > 45:
            posture = 'turned_back'
        else:
            posture = 'normal'
        
        avg_conf = np.mean([kp.confidence for kp in frame.keypoints 
                           if kp.confidence > 0])
        
        return {
            'posture': posture,
            'description': self.OOP_TYPES.get(posture, '未知'),
            'confidence': round(avg_conf, 3),
            'details': details
        }
    
    def analyze_sequence(self, frames: List[PoseFrame]) -> dict:
        """分析姿态序列"""
        results = [self.detect_posture(f) for f in frames]
        
        # 统计各姿态占比
        posture_counts = {}
        for r in results:
            p = r['posture']
            posture_counts[p] = posture_counts.get(p, 0) + 1
        
        # 主要姿态
        total = len(results)
        posture_pct = {k: round(v/total*100, 1) 
                       for k, v in posture_counts.items()}
        main_posture = max(posture_counts, key=posture_counts.get)
        
        # 姿态变化次数
        changes = sum(1 for i in range(1, len(results))
                     if results[i]['posture'] != results[i-1]['posture'])
        
        return {
            'total_frames': total,
            'posture_distribution': posture_pct,
            'main_posture': main_posture,
            'posture_changes': changes,
            'stability_score': round(100 - changes / max(total, 1) * 100, 1),
            'per_frame': results
        }


# ==================== 测试 ====================
if __name__ == "__main__":
    np.random.seed(42)
    detector = OOPDetector(640, 480)
    
    print("=" * 70)
    print("IMS OOP 异常姿态检测测试")
    print("参考: arXiv:2608.19480 VideoRun2D")
    print("=" * 70)
    
    # 模拟不同姿态
    def make_frame(nose, ls, rs, lh, rh, lk=None, rk=None, la=None, ra=None):
        kps = [
            Keypoint('nose', *nose[:2], nose[2] if len(nose)>2 else 0.9),
            Keypoint('left_shoulder', *ls[:2], 0.85),
            Keypoint('right_shoulder', *rs[:2], 0.85),
            Keypoint('left_hip', *lh[:2], 0.80),
            Keypoint('right_hip', *rh[:2], 0.80),
        ]
        if lk: kps.append(Keypoint('left_knee', *lk[:2], 0.75))
        if rk: kps.append(Keypoint('right_knee', *rk[:2], 0.75))
        return PoseFrame(kps, 0)
    
    scenarios = [
        ('正常坐姿', make_frame(
            (320, 100), (280, 180), (360, 180),
            (290, 320), (350, 320), (270, 420), (370, 420))),
        ('前倾看手机', make_frame(
            (340, 150), (290, 200), (360, 200),
            (295, 330), (355, 330), (270, 430), (370, 430))),
        ('侧倾靠窗', make_frame(
            (250, 110), (230, 190), (310, 170),
            (250, 330), (320, 310), (230, 430), (320, 410))),
        ('瘫坐', make_frame(
            (330, 120), (270, 190), (360, 190),
            (290, 290), (360, 290), (270, 410), (370, 410))),
        ('脚搭仪表台', make_frame(
            (320, 100), (280, 180), (360, 180),
            (290, 320), (350, 320), (250, 100), (370, 100))),
    ]
    
    print(f"\n{'场景':<15} {'姿态':>20} {'躯干角':>8} {'头部角':>8} "
          f"{'肩倾斜':>8} {'膝高度':>8} {'置信度':>8}")
    print("-" * 80)
    
    for name, frame in scenarios:
        r = detector.detect_posture(frame)
        d = r['details']
        print(f"{name:<15} {r['description']:>20} "
              f"{d['torso_angle']:>8.1f} {d['head_angle']:>8.1f} "
              f"{d['shoulder_tilt']:>8.1f} {d['knee_height_ratio']:>8.2f} "
              f"{r['confidence']:>8.3f}")
    
    # 序列分析
    print(f"\n{'='*70}")
    print("序列分析 (模拟30帧驾驶过程)")
    print(f"{'='*70}")
    
    seq = []
    for i in range(30):
        if i < 10:
            seq.append(scenarios[0][1])  # 正常
        elif i < 15:
            seq.append(scenarios[1][1])  # 前倾
        elif i < 20:
            seq.append(scenarios[3][1])  # 瘫坐
        else:
            seq.append(scenarios[0][1])  # 恢复正常
    
    analysis = detector.analyze_sequence(seq)
    print(f"\n  总帧数: {analysis['total_frames']}")
    print(f"  主要姿态: {analysis['main_posture']}")
    print(f"  姿态变化次数: {analysis['posture_changes']}")
    print(f"  稳定性评分: {analysis['stability_score']}/100")
    print(f"\n  姿态分布:")
    for p, pct in analysis['posture_distribution'].items():
        bar = '█' * int(pct / 2)
        print(f"    {p:<20} {pct:>6.1f}% {bar}")