import torch
import torch.nn as nn
import torchvision.models as models

class BiFuseNet(nn.Module):
    """
    BiFuseNet: 双流 3D 时空多模态网络
    
    用于从 RGB 面部视频估计酒精损伤等级
    
    架构：
    - 空间流：2D CNN 提取每帧面部特征
    - 时序流：3D CNN 捕获动态变化
    - 融合：双流特征拼接后分类
    
    论文核心思路：
    酒精损伤在面部的表现包括：
    1. 面部潮红（血管扩张）→ 空间特征
    2. 眼睑下垂、微表情异常 → 时空特征
    3. 头部晃动加剧 → 时序运动特征
    4. 眨眼频率变化 → 时序特征
    """
    
    def __init__(self, num_classes: int = 3, backbone: str = "resnet18"):
        super().__init__()
        
        # 空间流：2D ResNet 处理每帧
        self.spatial_backbone = models.resnet18(weights=models.ResNet18_Weights.DEFAULT)
        self.spatial_backbone.fc = nn.Identity()  # 移除分类头
        spatial_feat_dim = 512
        
        # 时序流：3D CNN（简化版）
        self.temporal_backbone = nn.Sequential(
            # 输入: (B, 3, T, H, W)
            nn.Conv3d(3, 32, kernel_size=(3, 3, 3), padding=(1, 1, 1)),
            nn.BatchNorm3d(32),
            nn.ReLU(),
            nn.MaxPool3d((1, 2, 2)),
            
            nn.Conv3d(32, 64, kernel_size=(3, 3, 3), padding=(1, 1, 1)),
            nn.BatchNorm3d(64),
            nn.ReLU(),
            nn.MaxPool3d((1, 2, 2)),
            
            nn.Conv3d(64, 128, kernel_size=(3, 3, 3), padding=(1, 1, 1)),
            nn.BatchNorm3d(128),
            nn.ReLU(),
            nn.AdaptiveAvgPool3d((1, 1, 1))
        )
        temporal_feat_dim = 128
        
        # 融合分类头
        fused_dim = spatial_feat_dim + temporal_feat_dim
        self.classifier = nn.Sequential(
            nn.Linear(fused_dim, 256),
            nn.ReLU(),
            nn.Dropout(0.5),
            nn.Linear(256, 128),
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(128, num_classes)
        )
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Args:
            x: 面部视频片段, shape=(B, C, T, H, W)
               B: batch size
               C: 3 (RGB)
               T: 帧数（论文使用 16 帧片段）
               H, W: 224x224
        
        Returns:
            logits: (B, 3) 三分类 logits
            [sober, low_AII, severe_AII]
        """
        B, C, T, H, W = x.shape
        
        # 空间流：逐帧提取 2D 特征
        x_2d = x.permute(0, 2, 1, 3, 4)  # (B, T, C, H, W)
        x_2d = x_2d.reshape(B * T, C, H, W)
        spatial_feat = self.spatial_backbone(x_2d)  # (B*T, 512)
        spatial_feat = spatial_feat.reshape(B, T, -1)
        spatial_feat = spatial_feat.mean(dim=1)  # 时间平均 → (B, 512)
        
        # 时序流：3D CNN
        temporal_feat = self.temporal_backbone(x)  # (B, 128, 1, 1, 1)
        temporal_feat = temporal_feat.flatten(1)  # (B, 128)
        
        # 融合
        fused = torch.cat([spatial_feat, temporal_feat], dim=-1)
        logits = self.classifier(fused)
        
        return logits


# 测试代码
if __name__ == "__main__":
    model = BiFuseNet(num_classes=3)
    
    # 模拟输入：16帧 224x224 RGB 面部视频
    x = torch.randn(4, 3, 16, 224, 224)
    
    with torch.no_grad():
        logits = model(x)
    
    print("=== BiFuseNet 测试 ===")
    print(f"输入: {x.shape}")
    print(f"输出: {logits.shape}")
    
    # 分类结果
    classes = ["sober", "low_AII", "severe_AII"]
    probs = torch.softmax(logits, dim=-1)
    for i, cls in enumerate(classes):
        print(f"  {cls}: {probs[0, i]:.4f}")
    
    # 参数量
    total_params = sum(p.numel() for p in model.parameters())
    print(f"\n参数量: {total_params:,}")
    print(f"模型大小: {total_params * 4 / 1024 / 1024:.1f} MB")

import numpy as np
from typing import List, Tuple

class AlcoholDatasetConfig:
    """
    BiFuseNet 数据集配置
    
    60 名受试者在模拟驾驶器中完成驾驶任务
    分三组采集面部视频：
    - sober: 0.00 g/dL (清醒)
    - low_AII: 0.03-0.05 g/dL (轻度损伤)  
    - severe_AII: >0.08 g/dL (重度损伤)
    
    数据采集协议：
    1. 受试者饮酒至目标 BAC
    2. 使用呼吸式酒精检测仪标定 BAC
    3. 在驾驶模拟器中驾驶 15 分钟
    4. RGB 摄像头 (1080p, 30fps) 录制面部
    5. 切分为 16 帧片段用于训练
    """
    
    SUBJECTS = 60
    CLASSES = ["sober", "low_AII", "severe_AII"]
    FRAMES_PER_CLIP = 16
    IMAGE_SIZE = (224, 224)
    FPS = 30
    
    # BAC 分级标准
    BAC_THRESHOLDS = {
        "sober": (0.00, 0.02),       # 0.00-0.02 g/dL
        "low_AII": (0.03, 0.05),     # 0.03-0.05 g/dL (WHO推荐上限)
        "severe_AII": (0.08, 0.15),  # ≥0.08 g/dL (多数国家法定上限)
    }
    
    # 数据增强策略
    AUGMENTATION = {
        "horizontal_flip": 0.5,
        "color_jitter": {"brightness": 0.2, "contrast": 0.2, "saturation": 0.1},
        "random_crop": True,
        "temporal_jitter": True,  # 帧序随机偏移
    }

class FacialAlcoholCues:
    """
    酒精损伤在面部的生理表现
    
    论文识别的关键面部线索：
    """
    
    CUES = {
        "vascular": {
            "description": "面部潮红（血管扩张）",
            "regions": ["脸颊", "鼻子", "额头"],
            "manifestation": "皮肤红度增加，局部温度升高",
            "detectability": "RGB 可见，但受肤色影响",
        },
        "ocular": {
            "description": "眼部异常",
            "features": [
                "眼睑下垂 (ptosis)",
                "眨眼频率变化",
                "凝视稳定性下降",
                "瞳孔对光反应迟缓"
            ],
            "detectability": "需要 ≥30fps 摄像头和高分辨率眼部区域",
        },
        "motor": {
            "description": "面部肌肉运动异常",
            "features": [
                "微表情减少（面具脸）",
                "嘴部不对称",
                "面部松弛度增加"
            ],
            "detectability": "需要时序建模（3D CNN 或 LSTM）",
        },
        "postural": {
            "description": "头部姿势变化",
            "features": [
                "头部晃动幅度增大",
                "前倾角度增加",
                "姿势稳定性下降"
            ],
            "detectability": "可通过面部关键点追踪检测",
        }
    }
    
    # 各特征的检测难度
    DETECTION_DIFFICULTY = {
        "面部潮红": {"rgb": "中等", "ir": "困难", "multi_modal": "中等"},
        "眼睑下垂": {"rgb": "容易", "ir": "容易", "multi_modal": "容易"},
        "眨眼频率": {"rgb": "中等", "ir": "中等", "multi_modal": "容易"},
        "微表情异常": {"rgb": "困难", "ir": "困难", "multi_modal": "中等"},
        "头部晃动": {"rgb": "容易", "ir": "容易", "multi_modal": "容易"},
    }

class MultiModalAlcoholDetection(nn.Module):
    """
    多模态酒驾检测融合方案
    
    结合 BiFuseNet (视觉) + mmWave 雷达 (生理)
    提升 BAC 估计精度和鲁棒性
    
    硬件配置：
    | 组件 | 型号 | 参数 | 安装位置 |
    |------|------|------|----------|
    | RGB-IR 摄像头 | OV2311 | 2MP, 1600x1200, 全局快门 | A柱/仪表盘上方 |
    | mmWave 雷达 | NOVELIC 60GHz | 57.7-62.4 GHz, 115° FoV | 仪表盘上方 |
    | 处理器 | QCS8255 | Hexagon NPU, 26 TOPS | 中控 |
    """
    
    def __init__(self):
        super().__init__()
        # 视觉分支
        self.visual_branch = BiFuseNet(num_classes=3)
        
        # 雷达分支（心率变异性 → BAC 间接推断）
        self.radar_branch = nn.Sequential(
            nn.Linear(6, 64),  # 输入: [HR, HRV, RR, RR_var, movement, HR_trend]
            nn.ReLU(),
            nn.Linear(64, 32),
            nn.ReLU(),
            nn.Linear(32, 3)
        )
        
        # 融合权重（可学习）
        self.fusion_weight = nn.Parameter(torch.tensor([0.7, 0.3]))  # 视觉70%, 雷达30%
        
    def forward(self, video: torch.Tensor, radar_features: torch.Tensor) -> dict:
        """
        Args:
            video: (B, 3, T, H, W) 面部视频
            radar_features: (B, 6) [HR, HRV, RR, RR_var, movement, HR_trend]
        
        Returns:
            output: {
                fused_logits: (B, 3) 融合分类,
                visual_logits: (B, 3) 视觉分类,
                radar_logits: (B, 3) 雷达分类,
                confidence: (B,) 置信度
            }
        """
        visual_logits = self.visual_branch(video)
        radar_logits = self.radar_branch(radar_features)
        
        # 加权融合
        w = torch.softmax(self.fusion_weight, dim=0)
        fused_logits = w[0] * visual_logits + w[1] * radar_logits
        
        return {
            "fused_logits": fused_logits,
            "visual_logits": visual_logits,
            "radar_logits": radar_logits,
            "confidence": torch.softmax(fused_logits, dim=-1).max(dim=-1)[0],
        }