import torch
import torch.nn as nn

class OccupantPostureDetector(nn.Module):
    """
    座舱乘员3D姿态检测器
    
    基于3D深度摄像头输入，输出17个身体关键点的3D坐标
    用于OOP判定
    """
    
    # COCO-17 关键点
    KEYPOINTS = [
        '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'
    ]
    
    def __init__(self, num_keypoints=17, num_posture_classes=6):
        super().__init__()
        
        # 骨干网络 (深度图输入)
        self.backbone = nn.Sequential(
            nn.Conv2d(1, 32, 3, stride=2, padding=1),  # 1通道深度图
            nn.BatchNorm2d(32), nn.ReLU(),
            nn.Conv2d(32, 64, 3, stride=2, padding=1),
            nn.BatchNorm2d(64), nn.ReLU(),
            nn.Conv2d(64, 128, 3, stride=2, padding=1),
            nn.BatchNorm2d(128), nn.ReLU(),
            nn.Conv2d(128, 256, 3, stride=2, padding=1),
            nn.BatchNorm2d(256), nn.ReLU(),
            nn.AdaptiveAvgPool2d((4, 4))
        )
        
        # 3D关键点回归头
        self.kp_head = nn.Sequential(
            nn.Flatten(),
            nn.Linear(256 * 4 * 4, 512),
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(512, num_keypoints * 3)  # 每个关键点 xyz
        )
        
        # 姿态分类头
        self.posture_head = nn.Sequential(
            nn.Linear(256 * 4 * 4, 256),
            nn.ReLU(),
            nn.Linear(256, num_posture_classes)
        )
        
        # OOP风险评分头
        self.risk_head = nn.Sequential(
            nn.Linear(num_keypoints * 3, 128),
            nn.ReLU(),
            nn.Linear(128, 1),
            nn.Sigmoid()
        )
    
    def forward(self, depth_map):
        """
        Args:
            depth_map: (B, 1, H, W) 深度图, 单位mm
            
        Returns:
            keypoints_3d: (B, 17, 3) 3D关键点坐标
            posture: (B, 6) 姿态分类logits
            risk: (B, 1) OOP风险评分 0-1
        """
        feat = self.backbone(depth_map)
        flat = feat.flatten(1)
        
        kp_3d = self.kp_head(flat).view(-1, 17, 3)
        posture = self.posture_head(flat)
        risk = self.risk_head(kp_3d.view(kp_3d.size(0), -1))
        
        return kp_3d, posture, risk


# 姿态类别
POSTURES = {
    0: '正常坐姿',
    1: '瘫坐',
    2: '前倾',
    3: '侧靠',
    4: '脚上仪表台',
    5: '儿童/小体型'
}

# OOP判定规则
OOP_RULES = {
    'head_to_dashboard_mm': 750,   # 头部距仪表台<750mm → OOP
    'torso_angle_deg': 30,          # 躯干前倾>30° → OOP
    'lateral_lean_mm': 200,         # 侧向偏移>200mm → OOP
    'feet_on_dash': True,           # 脚在仪表台 → OOP
    'child_size': True,             # 儿童体型 → 关闭气囊
}


def assess_oop(keypoints_3d, posture_logits, risk_score):
    """
    基于关键点和姿态评估OOP风险
    
    Returns:
        dict: OOP评估结果+约束系统建议
    """
    kp = keypoints_3d  # (17, 3)
    
    # 头部到仪表台距离 (近似: 鼻子z坐标)
    head_dist = kp[0, 2]  # nose z
    
    # 躯干角度 (肩膀中点到髋部中点)
    shoulder_mid = (kp[5] + kp[6]) / 2
    hip_mid = (kp[11] + kp[12]) / 2
    torso_vec = shoulder_mid - hip_mid
    torso_angle = torch.atan2(torso_vec[2], abs(torso_vec[1]) + 1e-6)
    torso_angle_deg = torch.rad2deg(torso_angle).item()
    
    # 侧向偏移
    lateral_offset = abs(kp[0, 0]).item()  # nose x
    
    # 判定
    is_oop = (
        head_dist < OOP_RULES['head_to_dashboard_mm'] or
        abs(torso_angle_deg) > OOP_RULES['torso_angle_deg'] or
        lateral_offset > OOP_RULES['lateral_lean_mm'] or
        risk_score > 0.7
    )
    
    # 约束系统建议
    if is_oop:
        if risk_score > 0.9:
            action = 'DEPLOY_LOW_POWER'  # 低功率展开
        elif risk_score > 0.7:
            action = 'DELAY_DEPLOY'  # 延迟展开
        else:
            action = 'WARN_ONLY'  # 仅告警
    else:
        action = 'DEPLOY_NORMAL'
    
    return {
        'is_oop': is_oop,
        'head_distance_mm': head_dist,
        'torso_angle_deg': torso_angle_deg,
        'lateral_offset_mm': lateral_offset,
        'risk_score': risk_score,
        'airbag_action': action
    }


if __name__ == "__main__":
    model = OccupantPostureDetector()
    
    # 模拟深度图输入 (480x640)
    depth = torch.randn(2, 1, 480, 640)
    
    kp, posture, risk = model(depth)
    
    print("=== OOP姿态检测 ===")
    print(f"输入: 深度图{depth.shape}")
    print(f"3D关键点: {kp.shape} (17点×3D)")
    print(f"姿态分类: {posture.shape}")
    print(f"风险评分: {risk.flatten().tolist()}")
    
    # 评估第一个乘员
    result = assess_oop(kp[0], posture[0], risk[0].item())
    print(f"\n乘员1评估:")
    print(f"  OOP: {'是 ⚠️' if result['is_oop'] else '否'}")
    print(f"  头部距离: {result['head_distance_mm']:.0f}mm")
    print(f"  躯干角度: {result['torso_angle_deg']:.1f}°")
    print(f"  侧向偏移: {result['lateral_offset_mm']:.0f}mm")
    print(f"  气囊动作: {result['airbag_action']}")
    
    print("\n=== 部署参数 ===")
    print("传感器: 3D深度摄像头 (TOF/结构光)")
    print("关键点: 17 (COCO格式)")
    print("姿态类别: 6类")
    print("延迟: <50ms (边缘NPU)")
    print("精度目标: 姿态分类>90%, 关键点误差<30mm")