import torch
import torch.nn as nn
import torch.nn.functional as F

class DrowSFormerXAI(nn.Module):
    """
    DrowSFormer-XAI: 跨注意力CNN-Transformer融合疲劳检测框架
    
    论文: Neural Processing Letters, 2026
    核心: CNN特征提取 + Transformer跨注意力融合 + XAI可解释性
    
    输入: 驾驶员面部视频帧序列
    输出: 疲劳等级 + 注意力热图(可解释性)
    """
    
    def __init__(self, num_classes=3, feat_dim=256, num_heads=8, 
                 num_transformer_layers=4):
        super().__init__()
        
        # 1. CNN特征提取（多流）
        self.eye_stream = FeatureStream(in_channels=3, feat_dim=feat_dim)
        self.mouth_stream = FeatureStream(in_channels=3, feat_dim=feat_dim)
        self.head_pose_stream = FeatureStream(in_channels=3, feat_dim=feat_dim)
        
        # 2. 跨注意力Transformer
        self.cross_attn = CrossAttentionFusion(
            feat_dim=feat_dim,
            num_heads=num_heads,
            num_layers=num_transformer_layers
        )
        
        # 3. 分类头
        self.classifier = nn.Sequential(
            nn.Linear(feat_dim, feat_dim // 2),
            nn.GELU(),
            nn.Dropout(0.1),
            nn.Linear(feat_dim // 2, num_classes)
        )
        
        # 4. 可解释性模块
        self.xai_module = XAIModule()
    
    def forward(self, eye_img, mouth_img, head_pose_img, return_attention=False):
        """
        Args:
            eye_img: (B, C, H, W) 眼部裁剪图像
            mouth_img: (B, C, H, W) 嘴部裁剪图像
            head_pose_img: (B, C, H, W) 头部姿态图像
            return_attention: 是否返回注意力图（可解释性）
        
        Returns:
            logits: (B, num_classes) 疲劳等级
            attention_maps: (可选) 注意力热图
        """
        # 1. 多流特征提取
        eye_feat = self.eye_stream(eye_img)       # (B, feat_dim)
        mouth_feat = self.mouth_stream(mouth_img) # (B, feat_dim)
        pose_feat = self.head_pose_stream(head_pose_img) # (B, feat_dim)
        
        # 2. 跨注意力融合
        fused, attn_weights = self.cross_attn(
            eye_feat, mouth_feat, pose_feat
        )
        
        # 3. 分类
        logits = self.classifier(fused)
        
        # 4. 可解释性
        if return_attention:
            attention_maps = self.xai_module.generate_attention_map(
                attn_weights, eye_img, mouth_img, head_pose_img
            )
            return logits, attention_maps
        
        return logits


class FeatureStream(nn.Module):
    """单流CNN特征提取器"""
    def __init__(self, in_channels=3, feat_dim=256):
        super().__init__()
        self.backbone = nn.Sequential(
            nn.Conv2d(in_channels, 32, 3, 1, 1),
            nn.BatchNorm2d(32),
            nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Conv2d(32, 64, 3, 1, 1),
            nn.BatchNorm2d(64),
            nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Conv2d(64, 128, 3, 1, 1),
            nn.BatchNorm2d(128),
            nn.ReLU(),
            nn.AdaptiveAvgPool2d(1)
        )
        self.proj = nn.Linear(128, feat_dim)
    
    def forward(self, x):
        feat = self.backbone(x).flatten(1)
        return self.proj(feat)


class CrossAttentionFusion(nn.Module):
    """
    跨注意力融合模块
    
    三个特征流之间的双向交叉注意力
    """
    def __init__(self, feat_dim=256, num_heads=8, num_layers=4):
        super().__init__()
        self.layers = nn.ModuleList([
            CrossAttentionLayer(feat_dim, num_heads) 
            for _ in range(num_layers)
        ])
        self.norm = nn.LayerNorm(feat_dim)
    
    def forward(self, eye, mouth, pose):
        """
        Args:
            eye, mouth, pose: (B, feat_dim)
        Returns:
            fused: (B, feat_dim)
            attn_weights: 注意力权重（用于可解释性）
        """
        # 堆叠为序列: (B, 3, feat_dim)
        x = torch.stack([eye, mouth, pose], dim=1)
        
        all_attn = []
        for layer in self.layers:
            x, attn = layer(x)
            all_attn.append(attn)
        
        # 融合序列
        fused = self.norm(x.mean(dim=1))
        return fused, all_attn


class CrossAttentionLayer(nn.Module):
    """跨注意力层"""
    def __init__(self, dim=256, num_heads=8):
        super().__init__()
        self.attn = nn.MultiheadAttention(dim, num_heads, batch_first=True)
        self.norm1 = nn.LayerNorm(dim)
        self.norm2 = nn.LayerNorm(dim)
        self.ffn = nn.Sequential(
            nn.Linear(dim, dim * 4),
            nn.GELU(),
            nn.Linear(dim * 4, dim)
        )
    
    def forward(self, x):
        attn_out, attn_weights = self.attn(x, x, x)
        x = self.norm1(x + attn_out)
        x = self.norm2(x + self.ffn(x))
        return x, attn_weights


class XAIModule:
    """
    可解释性模块
    
    生成注意力热图和特征归因
    """
    def generate_attention_map(self, attn_weights, eye_img, mouth_img, pose_img):
        """
        从Transformer注意力权重生成可视化热图
        
        Returns:
            dict: 各模态的注意力权重
        """
        # 最后层的注意力: (B, 3, 3)
        last_attn = attn_weights[-1]
        
        # 三个流的注意力权重
        eye_weight = last_attn[:, 0].mean(dim=-1)    # 眼部关注度
        mouth_weight = last_attn[:, 1].mean(dim=-1)  # 嘴部关注度
        pose_weight = last_attn[:, 2].mean(dim=-1)   # 头部姿态关注度
        
        return {
            'eye_attention': eye_weight,
            'mouth_attention': mouth_weight,
            'pose_attention': pose_weight,
            'fatigue_indicators': {
                'eye_closure': '眼部闭眼频率',
                'yawning': '嘴部张开程度',
                'head_nodding': '头部点头频率'
            }
        }


# 测试
if __name__ == "__main__":
    model = DrowSFormerXAI(num_classes=3, feat_dim=256)
    
    # 模拟输入
    eye = torch.randn(4, 3, 64, 64)
    mouth = torch.randn(4, 3, 64, 64)
    pose = torch.randn(4, 3, 64, 64)
    
    logits, attn = model(eye, mouth, pose, return_attention=True)
    
    print(f"Logits: {logits.shape}")  # (4, 3): 清醒/轻度疲劳/严重疲劳
    print(f"Eye attention: {attn['eye_attention']}")
    print(f"Mouth attention: {attn['mouth_attention']}")
    print(f"Pose attention: {attn['pose_attention']}")
    print(f"Parameters: {sum(p.numel() for p in model.parameters()):,}")

# IMS集成方案
class IMSDrowsinessModule:
    def __init__(self):
        self.drowsformer = DrowSFormerXAI()
        self.perclos_tracker = PERCLOSTracker()
        
    def detect(self, face_frames):
        # 方法1: DrowSFormer-XAI（深度学习）
        dl_result, attn = self.drowsformer(face_frames, return_attention=True)
        
        # 方法2: PERCLOS（规则方法）
        perclos_result = self.perclos_tracker(face_frames)
        
        # 融合决策
        if dl_result.confidence > 0.85:
            return dl_result, attn  # 用可解释结果
        else:
            return self._fuse(dl_result, perclos_result)