"""
安全带检测与误用分类网络

目标: 从车内摄像头图像检测安全带状态
分类: 6类 (正常/未系/身后/腋下/未锁扣/儿童干扰)

模型: 轻量级目标检测 + 关键点估计
"""

import torch
import torch.nn as nn

class SeatbeltDetector(nn.Module):
    """
    安全带检测器
    
    架构: 
    1. 骨干网络: MobileNetV3-Small (轻量)
    2. 检测头: 安全带区域检测
    3. 关键点头: 锚点+锁扣位置
    4. 分类头: 6类状态分类
    
    总参数: ~3M, INT8后 ~800KB
    """
    
    def __init__(self, num_classes=6):
        super().__init__()
        
        # 轻量骨干
        self.backbone = MobileNetV3Small(in_channels=3)
        
        # 安全带区域检测头
        self.belt_det = nn.Sequential(
            nn.Conv2d(128, 64, 1),
            nn.ReLU(),
            nn.Conv2d(64, 2, 1)  # 2 anchors
        )
        
        # 关键点检测头 (肩部锚点, 锁扣位置, 腰带位置)
        self.keypoints = nn.Sequential(
            nn.Conv2d(128, 64, 1),
            nn.ReLU(),
            nn.Conv2d(64, 6, 1)  # 3个关键点 × 2(x,y)
        )
        
        # 分类头
        self.classifier = nn.Sequential(
            nn.AdaptiveAvgPool2d(1),
            nn.Flatten(),
            nn.Linear(128, 64),
            nn.ReLU(),
            nn.Dropout(0.2),
            nn.Linear(64, num_classes)
        )
    
    def forward(self, x):
        features = self.backbone(x)
        belt_boxes = self.belt_det(features)
        belt_keypoints = self.keypoints(features)
        logits = self.classifier(features)
        return logits, belt_boxes, belt_keypoints


class MobileNetV3Small(nn.Module):
    """轻量骨干网络"""
    def __init__(self, in_channels=3):
        super().__init__()
        self.features = nn.Sequential(
            nn.Conv2d(in_channels, 16, 3, 2, 1, bias=False),
            nn.BatchNorm2d(16),
            nn.Hardswish(inplace=True),
            # 简化的MobileNetV3块
            nn.Conv2d(16, 32, 3, 2, 1, groups=1, bias=False),
            nn.BatchNorm2d(32),
            nn.ReLU6(inplace=True),
            nn.Conv2d(32, 64, 3, 2, 1, groups=1, bias=False),
            nn.BatchNorm2d(64),
            nn.ReLU6(inplace=True),
            nn.Conv2d(64, 128, 3, 2, 1, bias=False),
            nn.BatchNorm2d(128),
            nn.Hardswish(inplace=True),
        )
    
    def forward(self, x):
        return self.features(x)


# ============================================
# 误用判断逻辑
# ============================================

class SeatbeltMisuseClassifier:
    """
    基于关键点几何关系判断安全带误用
    
    关键点定义:
    - kp0: 肩带锚点 (B柱上方)
    - kp1: 锁扣位置
    - kp2: 腰带位置 (骨盆附近)
    """
    
    def classify(self, keypoints, person_pose):
        """
        Args:
            keypoints: (3, 2) 三个关键点坐标
            person_pose: dict with shoulders, hips, chest
        Returns:
            category: str, confidence: float
        """
        anchor, buckle, lap = keypoints
        
        # 1. 未系: 无安全带特征
        if self._belt_not_visible(keypoints):
            return 'unbelted', 0.95
        
        # 2. 正常佩戴: 肩带从锚点→跨胸前→锁扣
        if self._normal_pattern(keypoints, person_pose):
            return 'normal', 0.90
        
        # 3. 身后佩戴: 肩带方向背离胸前
        if self._behind_back_pattern(keypoints, person_pose):
            return 'behind_back', 0.75
        
        # 4. 腋下佩戴: 肩带在腋下而非肩上
        if self._underarm_pattern(keypoints, person_pose):
            return 'underarm', 0.70
        
        # 5. 未锁扣: 安全带拉出但扣环未闭合
        if self._unbuckled_pattern(keypoints):
            return 'unbuckled', 0.65
        
        return 'uncertain', 0.3
    
    def _belt_not_visible(self, kps):
        """安全带不可见"""
        return all(k[0] < 0 for k in kps)
    
    def _normal_pattern(self, kps, pose):
        """正常佩戴: 锚点→胸前→锁扣"""
        anchor, buckle, lap = kps
        chest = pose.get('chest', (0.5, 0.4))
        return (anchor[1] < chest[1] < buckle[1] and  # 锚点在上
                abs(anchor[0] - buckle[0]) < 0.3)  # 水平距离合理
    
    def _behind_back_pattern(self, kps, pose):
        """身后佩戴"""
        anchor, buckle, lap = kps
        # 肩带方向异常: 不经过胸前
        return anchor[1] < buckle[1] and abs(anchor[0] - buckle[0]) > 0.4
    
    def _underarm_pattern(self, kps, pose):
        """腋下佩戴"""
        anchor, buckle, lap = kps
        armpit = pose.get('right_armpit', (0.6, 0.5))
        # 肩带经过腋下而非肩上
        return abs(anchor[0] - armpit[0]) < 0.1
    
    def _unbuckled_pattern(self, kps):
        """未锁扣"""
        anchor, buckle, lap = kps
        # 锁扣位置特征异常
        return buckle[0] < 0 or buckle[1] < 0