import torch
import torch.nn as nn
from typing import Dict, List, Tuple

class OccupantClassifier(nn.Module):
    """
    Aptiv风格的乘员分类模型
    
    基于2D姿态估计+尺寸分析
    输入：内舱摄像头帧
    输出：乘员类别+置信度
    
    类别：adult, child_fwd, child_rear, infant, object, empty
    """
    
    def __init__(self, num_classes: int = 6):
        super().__init__()
        
        # 骨干网络（轻量化，适合边缘部署）
        self.backbone = nn.Sequential(
            # 简化版MobileNet
            nn.Conv2d(3, 32, 3, stride=2, padding=1),
            nn.BatchNorm2d(32),
            nn.ReLU6(inplace=True),
            
            # Depthwise Separable Conv
            nn.Conv2d(32, 32, 3, groups=32, stride=1, padding=1),
            nn.Conv2d(32, 64, 1),
            nn.BatchNorm2d(64),
            nn.ReLU6(inplace=True),
            
            nn.Conv2d(64, 64, 3, groups=64, stride=2, padding=1),
            nn.Conv2d(64, 128, 1),
            nn.BatchNorm2d(128),
            nn.ReLU6(inplace=True),
            
            nn.Conv2d(128, 128, 3, groups=128, stride=1, padding=1),
            nn.Conv2d(128, 256, 1),
            nn.BatchNorm2d(256),
            nn.ReLU6(inplace=True),
            
            nn.AdaptiveAvgPool2d(1)
        )
        
        # 分类头
        self.classifier = nn.Sequential(
            nn.Dropout(0.2),
            nn.Linear(256, 128),
            nn.ReLU6(inplace=True),
            nn.Linear(128, num_classes)
        )
        
        # 尺寸回归头（辅助任务）
        self.size_regressor = nn.Sequential(
            nn.Linear(256, 64),
            nn.ReLU6(inplace=True),
            nn.Linear(64, 4)  # (height_ratio, width_ratio, shoulder_w, head_size)
        )
        
        # CRS检测头（儿童座椅检测）
        self.crs_detector = nn.Sequential(
            nn.Linear(256, 64),
            nn.ReLU6(inplace=True),
            nn.Linear(64, 3)  # (no_crs, rear_facing, forward_facing)
        )
    
    def forward(self, x: torch.Tensor) -> Dict[str, torch.Tensor]:
        """
        Args:
            x: (B, 3, 224, 224) 座椅区域裁剪
        
        Returns:
            class_logits: (B, 6) 分类logits
            size_pred: (B, 4) 尺寸预测
            crs_pred: (B, 3) CRS检测
        """
        features = self.backbone(x)
        features = features.flatten(1)
        
        class_logits = self.classifier(features)
        size_pred = torch.sigmoid(self.size_regressor(features))
        crs_pred = self.crs_detector(features)
        
        return {
            'class_logits': class_logits,
            'size_pred': size_pred,
            'crs_pred': crs_pred
        }


class OccupantClassificationSystem:
    """
    完整乘员分类系统
    
    决策逻辑：
    1. 检测座椅是否有人
    2. 如有人，检测是否有CRS
    3. 如有CRS，判断前后向
    4. 如无CRS，根据尺寸判断成人/儿童
    5. 根据类别决定安全气囊策略
    """
    
    CLASS_NAMES = ['adult', 'child_fwd', 'child_rear', 
                   'infant', 'object', 'empty']
    
    AIRBAG_POLICY = {
        'adult': 'normal_deploy',
        'child_fwd': 'reduced_deploy',
        'child_rear': 'disable',
        'infant': 'disable',
        'object': 'disable',
        'empty': 'disable'
    }
    
    def __init__(self, confidence_threshold: float = 0.85):
        self.model = OccupantClassifier()
        self.threshold = confidence_threshold
        self.model.eval()
    
    def classify(self, seat_image: torch.Tensor) -> Dict:
        """
        分类单个座椅
        
        Args:
            seat_image: (1, 3, 224, 224) 座椅区域图像
        
        Returns:
            classification: 乘员类别+安全气囊策略
        """
        with torch.no_grad():
            output = self.model(seat_image)
        
        # 分类
        class_probs = torch.softmax(output['class_logits'], dim=-1)
        confidence, predicted = torch.max(class_probs, dim=-1)
        
        class_name = self.CLASS_NAMES[predicted.item()]
        conf = confidence.item()
        
        # CRS检测
        crs_probs = torch.softmax(output['crs_pred'], dim=-1)
        crs_type = ['no_crs', 'rear_facing', 'forward_facing'][
            torch.argmax(crs_probs, dim=-1).item()
        ]
        
        # 尺寸估计
        size = output['size_pred'][0].tolist()
        
        # 决策
        airbag = self.AIRBAG_POLICY.get(class_name, 'disable')
        
        # 置信度检查
        if conf < self.threshold:
            # 低置信度→保守策略
            airbag = 'disable'
            class_name = 'uncertain'
        
        return {
            'class': class_name,
            'confidence': conf,
            'crs_type': crs_type,
            'size': {
                'height_ratio': size[0],
                'width_ratio': size[1],
                'shoulder_width': size[2],
                'head_size': size[3]
            },
            'airbag_policy': airbag
        }


# 测试
if __name__ == "__main__":
    system = OccupantClassificationSystem(confidence_threshold=0.85)
    
    total_params = sum(p.numel() for p in system.model.parameters())
    print(f"模型参数量: {total_params:,} ({total_params/1e6:.2f}M)")
    
    # 模拟测试
    scenarios = {
        "成人正常坐姿": torch.randn(1, 3, 224, 224),
        "后向儿童座椅": torch.randn(1, 3, 224, 224) * 0.8,
        "空座椅": torch.randn(1, 3, 224, 224) * 0.3,
        "购物袋": torch.randn(1, 3, 224, 224) * 0.5,
    }
    
    print(f"\n{'场景':15s} | {'类别':10s} | {'置信度':>6} | {'CRS':>15} | {'气囊策略'}")
    print("-" * 70)
    
    for name, img in scenarios.items():
        result = system.classify(img)
        print(f"{name:15s} | {result['class']:10s} | "
              f"{result['confidence']:6.1%} | {result['crs_type']:>15} | "
              f"{result['airbag_policy']}")