import torch
import torch.nn as nn

class AlcoholImpairmentDetector(nn.Module):
    """
    纯视觉+行为酒精损伤检测
    
    基于 Greer et al. (2023) 多维分析思路
    Euro NCAP 2026 酒驾检测需求
    
    输入：
        - face_video: (B, T, 3, H, W) 面部视频
        - driving_signal: (B, T, 4) CAN信号 (speed, steer, lane_pos, accel)
    输出：
        - impairment_logits: (B, 3) 正常/轻度/重度损伤
    """
    
    def __init__(self, hidden=128):
        super().__init__()
        
        # 1. 面部特征编码器
        self.face_encoder = nn.Sequential(
            nn.Conv2d(3, 32, 3, stride=2, padding=1), nn.ReLU(),
            nn.Conv2d(32, 64, 3, stride=2, padding=1), nn.ReLU(),
            nn.Conv2d(64, 128, 3, stride=2, padding=1), nn.ReLU(),
            nn.AdaptiveAvgPool2d((2, 2)),
        )
        self.face_lstm = nn.LSTM(128*4, hidden, batch_first=True, num_layers=2, dropout=0.2)
        
        # 2. 驾驶行为编码器（CAN信号）
        self.drive_encoder = nn.Sequential(
            nn.Linear(4, 64), nn.ReLU(),
            nn.Linear(64, hidden), nn.ReLU(),
        )
        self.drive_lstm = nn.LSTM(hidden, hidden, batch_first=True)
        
        # 3. 融合+分类
        self.fusion = nn.Sequential(
            nn.Linear(hidden * 2, hidden), nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(hidden, 3)  # 正常/轻度/重度
        )
        
        # 4. 连续损伤程度回归（BAC估计）
        self.bac_regressor = nn.Sequential(
            nn.Linear(hidden * 2, 32), nn.ReLU(),
            nn.Linear(32, 1)  # BAC估计值
        )
    
    def forward(self, face_video, driving_signal):
        B, T, C, H, W = face_video.shape
        
        # 面部特征
        face_feat = self.face_encoder(face_video.view(B*T, C, H, W))
        face_feat = face_feat.view(B, T, -1)
        _, (face_hn, _) = self.face_lstm(face_feat)
        face_out = face_hn[-1]  # (B, hidden)
        
        # 驾驶行为
        drive_feat = self.drive_encoder(driving_signal)  # (B, T, hidden)
        _, (drive_hn, _) = self.drive_lstm(drive_feat)
        drive_out = drive_hn[-1]  # (B, hidden)
        
        # 融合
        combined = torch.cat([face_out, drive_out], dim=1)
        
        return {
            'impairment': self.fusion(combined),
            'bac_estimate': self.bac_regressor(combined),
        }


# 测试
if __name__ == "__main__":
    model = AlcoholImpairmentDetector()
    
    face = torch.randn(2, 60, 3, 64, 64)
    driving = torch.randn(2, 60, 4)
    
    out = model(face, driving)
    print(f"损伤分类: {out['impairment'].shape}")
    print(f"BAC估计: {out['bac_estimate'].shape}")