"""
GDAU (Guided Dual-Attention Unit) 架构复现

论文核心模块：
1. Position-Aware Spatial Attention（位置感知空间注意力）
2. SE Channel Attention（Squeeze-Excitation 通道注意力）

创新点：两种注意力串联使用，空间注意力先聚焦关键区域，
通道注意力再增强重要特征通道

对比评估的 13 种变体包括：
- 无注意力（基线）
- SE (Squeeze-Excitation)
- CBAM (Convolutional Block Attention Module)
- ECA (Efficient Channel Attention)
- CA (Coordinate Attention)
- DA (Dual Attention)
- GDAU (本文提出)
- 以及以上变体的组合
"""

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

class PositionAwareSpatialAttention(nn.Module):
    """
    位置感知空间注意力
    
    传统空间注意力只关注"哪里重要"，
    位置感知还编码"位置先验"——
    在 DMS 中，眼睛/嘴巴的位置有先验分布
    """
    
    def __init__(self, kernel_size: int = 7):
        super().__init__()
        # 位置编码
        self.pos_embed = nn.Parameter(torch.randn(1, 2, 1, 1))
        # 空间注意力
        self.conv = nn.Conv2d(2, 1, kernel_size, padding=kernel_size//2, bias=False)
        
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Args:
            x: (B, C, H, W)
        Returns:
            x_attended: (B, C, H, W)
        """
        B, C, H, W = x.shape
        
        # 位置编码广播到输入尺寸
        pos = F.interpolate(self.pos_embed, size=(H, W), mode='bilinear', align_corners=False)
        
        # 拼接位置信息到通道注意力
        avg_pool = x.mean(dim=1, keepdim=True)  # (B, 1, H, W)
        max_pool = x.max(dim=1, keepdim=True)[0]  # (B, 1, H, W)
        
        # 加入位置信息
        spatial = torch.cat([avg_pool + pos[:, 0:1], max_pool + pos[:, 1:2]], dim=1)
        attn = torch.sigmoid(self.conv(spatial))
        
        return x * attn


class SEChannelAttention(nn.Module):
    """
    Squeeze-Excitation 通道注意力
    
    压缩空间维度 → 学习通道权重 → 激活重要通道
    """
    
    def __init__(self, channels: int, reduction: int = 16):
        super().__init__()
        self.fc = nn.Sequential(
            nn.AdaptiveAvgPool2d(1),
            nn.Flatten(),
            nn.Linear(channels, channels // reduction),
            nn.ReLU(),
            nn.Linear(channels // reduction, channels),
            nn.Sigmoid()
        )
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        w = self.fc(x).unsqueeze(-1).unsqueeze(-1)
        return x * w


class GDAU(nn.Module):
    """
    Guided Dual-Attention Unit
    
    顺序：空间注意力 → 通道注意力
    "Guided" 指空间注意力引导通道注意力的聚焦方向
    """
    
    def __init__(self, channels: int, reduction: int = 16):
        super().__init__()
        self.spatial_attn = PositionAwareSpatialAttention()
        self.channel_attn = SEChannelAttention(channels, reduction)
        # 残差连接
        self.norm = nn.BatchNorm2d(channels)
        
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        residual = x
        # 空间注意力先聚焦
        x = self.spatial_attn(x)
        # 通道注意力增强
        x = self.channel_attn(x)
        # 残差
        return self.norm(residual + x)


class CompactDrowsinessNet(nn.Module):
    """
    紧凑疲劳检测网络
    
    基于 MobileNetV2 backbone + GDAU 模块
    
    设计目标：
    - 精度 ≥95%
    - 参数量 <2M
    - INT8 量化后 <2MB
    - FPGA 可部署
    """
    
    def __init__(self, num_classes: int = 2):
        super().__init__()
        
        # 轻量 backbone（MobileNetV2 精简版）
        self.backbone = nn.Sequential(
            # Entry
            nn.Conv2d(3, 32, 3, stride=2, padding=1, bias=False),
            nn.BatchNorm2d(32),
            nn.ReLU6(),
            
            # Block 1
            self._conv_bn(32, 16, 1),
            self._conv_dw(16, 32, 1),  # depthwise
            
            # GDAU 1
            GDAU(32, reduction=8),
            
            # Block 2
            self._conv_dw(32, 64, 2),
            self._conv_bn(64, 32, 1),
            
            # GDAU 2
            GDAU(32, reduction=8),
            
            # Block 3
            self._conv_dw(32, 128, 2),
            self._conv_bn(128, 64, 1),
            
            # GDAU 3
            GDAU(64, reduction=16),
            
            # Block 4
            self._conv_dw(64, 128, 1),
            
            # Head
            nn.AdaptiveAvgPool2d(1),
            nn.Flatten(),
            nn.Dropout(0.2),
            nn.Linear(128, num_classes)
        )
    
    def _conv_bn(self, inp, oup, stride):
        return nn.Sequential(
            nn.Conv2d(inp, oup, 3, stride, 1, bias=False),
            nn.BatchNorm2d(oup),
            nn.ReLU6()
        )
    
    def _conv_dw(self, inp, oup, stride):
        return nn.Sequential(
            nn.Conv2d(inp, inp, 3, stride, 1, groups=inp, bias=False),
            nn.BatchNorm2d(inp),
            nn.ReLU6(),
            nn.Conv2d(inp, oup, 1, 1, 0, bias=False),
            nn.BatchNorm2d(oup),
            nn.ReLU6()
        )
    
    def forward(self, x):
        return self.backbone(x)


# ==================== 测试代码 ====================
if __name__ == "__main__":
    import time
    
    model = CompactDrowsinessNet(num_classes=2)
    model.eval()
    
    # 输入：224x224 面部图像
    x = torch.randn(1, 3, 224, 224)
    
    # FP32 推理
    start = time.time()
    with torch.no_grad():
        out = model(x)
    fp32_ms = (time.time() - start) * 1000
    
    # 参数统计
    total_params = sum(p.numel() for p in model.parameters())
    
    print("=" * 60)
    print("CompactDrowsinessNet (GDAU) 测试")
    print("=" * 60)
    print(f"输入: {x.shape}")
    print(f"输出: {out.shape}")
    print(f"参数量: {total_params:,} ({total_params/1e6:.2f}M)")
    print(f"FP32 推理: {fp32_ms:.2f} ms ({1000/fp32_ms:.1f} FPS)")
    
    # 模拟 INT8 量化
    print(f"\n--- INT8 量化预估 ---")
    print(f"FP32 模型大小: {total_params * 4 / 1024:.1f} KB")
    print(f"INT8 模型大小: {total_params / 1024:.1f} KB")
    print(f"INT8 推理预估: {fp32_ms * 0.4:.2f} ms ({1000/(fp32_ms*0.4):.1f} FPS)")
    print(f"精度损失预估: 1-3%")
    
    # GDAU 模块分析
    gdau_count = sum(1 for m in model.modules() if isinstance(m, GDAU))
    print(f"\nGDAU 模块数: {gdau_count}")
    print(f"GDAU 参数占比: {sum(p.numel() for m in model.modules() if isinstance(m, GDAU) for p in m.parameters()) / total_params * 100:.1f}%")

"""
为什么校准重要？

场景：模型预测"疲劳"概率 0.95
- 校准好：95% 的情况下确实疲劳 → 可信告警
- 校准差：可能只有 70% 真实疲劳 → 大量误报

ECE (Expected Calibration Error):
- 衡量预测概率与真实频率的偏差
- GDAU ECE=2.1%：预测0.95时，真实疲劳频率~93%
- 基线 ECE=7.8%：预测0.95时，真实疲劳频率~87%

IMS 需要校准好的模型，因为告警阈值直接依赖预测概率
"""