"""
ECU BiFuseNet 统一驾驶损伤检测模型架构

论文: BMVC 2026
机构: Edith Cowan University

核心方法:
1. 3D 面部特征提取（同时处理时序面部序列）
2. 多任务分类头：疲劳 / 酒精 / 情绪
3. 双输入融合：RGB + IR 红外视频
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import Dict, Tuple

class BiFuseNetBackbone(nn.Module):
    """
    BiFuseNet 双流主干网络
    
    同时处理 RGB 和 IR（红外）视频流，
    在特征层进行融合，提升低光照条件下检测精度。
    
    论文描述：BiFuseNet 追踪细粒度面部变化：
    - 眨眼频率（blink rate）
    - 微小肌肉移动（small muscle shifts）
    - 面部不对称性（facial asymmetry）
    """
    
    def __init__(self, config: dict):
        super().__init__()
        
        # RGB 流：标准 3D CNN 提取时序面部特征
        self.rgb_stream = nn.Sequential(
            nn.Conv3d(3, 32, kernel_size=(3, 3, 3), padding=1),
            nn.BatchNorm3d(32),
            nn.ReLU(),
            nn.MaxPool3d(kernel_size=(1, 2, 2)),
            
            nn.Conv3d(32, 64, kernel_size=(3, 3, 3), padding=1),
            nn.BatchNorm3d(64),
            nn.ReLU(),
            nn.MaxPool3d(kernel_size=(1, 2, 2)),
            
            nn.Conv3d(64, 128, kernel_size=(3, 3, 3), padding=1),
            nn.BatchNorm3d(128),
            nn.ReLU(),
            nn.AdaptiveAvgPool3d((1, 1, 1))
        )
        
        # IR 流：红外视频流，结构与 RGB 流对称
        self.ir_stream = nn.Sequential(
            nn.Conv3d(1, 32, kernel_size=(3, 3, 3), padding=1),
            nn.BatchNorm3d(32),
            nn.ReLU(),
            nn.MaxPool3d(kernel_size=(1, 2, 2)),
            
            nn.Conv3d(32, 64, kernel_size=(3, 3, 3), padding=1),
            nn.BatchNorm3d(64),
            nn.ReLU(),
            nn.MaxPool3d(kernel_size=(1, 2, 2)),
            
            nn.Conv3d(64, 128, kernel_size=(3, 3, 3), padding=1),
            nn.BatchNorm3d(128),
            nn.ReLU(),
            nn.AdaptiveAvgPool3d((1, 1, 1))
        )
        
        # 融合层：将 RGB 和 IR 特征在通道维度拼接后融合
        fusion_dim = 128 + 128  # 双流各 128 维
        self.fusion = nn.Sequential(
            nn.Linear(fusion_dim, 256),
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(256, 128),
            nn.ReLU()
        )
        
    def forward(self, rgb: torch.Tensor, ir: torch.Tensor = None) -> torch.Tensor:
        """
        前向传播
        
        Args:
            rgb: RGB 视频片段, shape=(B, C=3, T, H, W)
            ir: IR 红外视频片段, shape=(B, C=1, T, H, W)（可选）
            
        Returns:
            fused_features: 融合后的 128 维特征向量
        """
        # RGB 流
        rgb_feat = self.rgb_stream(rgb)  # (B, 128, 1, 1, 1)
        rgb_feat = rgb_feat.flatten(1)  # (B, 128)
        
        if ir is not None:
            # IR 流
            ir_feat = self.ir_stream(ir)
            ir_feat = ir_feat.flatten(1)
            # 双流融合
            combined = torch.cat([rgb_feat, ir_feat], dim=1)
        else:
            # 仅 RGB 输入时，IR 特征用零向量替代
            combined = torch.cat([rgb_feat, torch.zeros_like(rgb_feat)], dim=1)
        
        fused = self.fusion(combined)
        return fused


class UnifiedImpairmentHead(nn.Module):
    """
    多任务损伤检测头
    
    从统一特征向量分别预测：
    - 疲劳等级（3级：清醒/轻度疲劳/严重疲劳）
    - 酒精损伤（2分类：清醒/受损）
    - 情绪状态（4分类：中性/愤怒/悲伤/焦虑）
    """
    
    def __init__(self, feature_dim: int = 128):
        super().__init__()
        
        # 疲劳检测头
        self.fatigue_head = nn.Sequential(
            nn.Linear(feature_dim, 64),
            nn.ReLU(),
            nn.Linear(64, 3)  # 清醒/轻度/严重
        )
        
        # 酒精损伤检测头
        self.alcohol_head = nn.Sequential(
            nn.Linear(feature_dim, 64),
            nn.ReLU(),
            nn.Linear(64, 2)  # 清醒/受损
        )
        
        # 情绪检测头
        self.emotion_head = nn.Sequential(
            nn.Linear(feature_dim, 64),
            nn.ReLU(),
            nn.Linear(64, 4)  # 中性/愤怒/悲伤/焦虑
        )
    
    def forward(self, features: torch.Tensor) -> Dict[str, torch.Tensor]:
        return {
            'fatigue': self.fatigue_head(features),
            'alcohol': self.alcohol_head(features),
            'emotion': self.emotion_head(features)
        }


class BiFuseNet(nn.Module):
    """
    完整的 BiFuseNet 统一驾驶损伤检测模型
    
    结合双流主干网络 + 多任务检测头
    
    性能（论文报告）:
    - 疲劳检测准确率: 95%
    - 酒精损伤准确率: 88.41%
    - 情绪检测: 持续优化中
    
    输入: RGB (B,3,T,H,W) + IR (B,1,T,H,W)
    输出: Dict[str, Tensor] 三种损伤状态预测
    """
    
    def __init__(self, config: dict = None):
        super().__init__()
        config = config or {}
        self.backbone = BiFuseNetBackbone(config)
        self.head = UnifiedImpairmentHead(feature_dim=128)
        
    def forward(self, rgb: torch.Tensor, ir: torch.Tensor = None) -> Dict[str, torch.Tensor]:
        features = self.backbone(rgb, ir)
        return self.head(features)


# ==================== 测试代码 ====================
if __name__ == "__main__":
    import numpy as np
    
    # 模型初始化
    model = BiFuseNet()
    model.eval()
    
    # 模拟输入：16帧 RGB + IR 视频片段
    batch_size = 1
    frames = 16
    height, width = 224, 224
    
    rgb_input = torch.randn(batch_size, 3, frames, height, width)
    ir_input = torch.randn(batch_size, 1, frames, height, width)
    
    # 推理
    with torch.no_grad():
        outputs = model(rgb_input, ir_input)
    
    # 结果解读
    fatigue_labels = ['清醒', '轻度疲劳', '严重疲劳']
    alcohol_labels = ['清醒', '酒精损伤']
    emotion_labels = ['中性', '愤怒', '悲伤', '焦虑']
    
    fatigue_pred = fatigue_labels[outputs['fatigue'].argmax(1).item()]
    alcohol_pred = alcohol_labels[outputs['alcohol'].argmax(1).item()]
    emotion_pred = emotion_labels[outputs['emotion'].argmax(1).item()]
    
    print("=" * 60)
    print("BiFuseNet 统一驾驶损伤检测结果")
    print("=" * 60)
    print(f"疲劳状态: {fatigue_pred}")
    print(f"酒精损伤: {alcohol_pred}")
    print(f"情绪状态: {emotion_pred}")
    print(f"\n疲劳置信度: {F.softmax(outputs['fatigue'], dim=1)[0]}")
    print(f"酒精置信度: {F.softmax(outputs['alcohol'], dim=1)[0]}")
    print(f"情绪置信度: {F.softmax(outputs['emotion'], dim=1)[0]}")
    
    # 模型参数统计
    total_params = sum(p.numel() for p in model.parameters())
    print(f"\n模型参数量: {total_params:,}")
    
    # 推理速度测试
    import time
    start = time.time()
    for _ in range(10):
        with torch.no_grad():
            _ = model(rgb_input, ir_input)
    avg_ms = (time.time() - start) / 10 * 1000
    print(f"平均推理时间: {avg_ms:.1f} ms ({1000/avg_ms:.1f} FPS)")