"""
DrowSFormer-XAI 核心实现
论文: Rajalingam et al., Neural Processing Letters 2026
DOI: 10.1007/s11063-026-11841-6

依赖: torch>=2.0, timm>=0.9, torchvision>=0.15
"""

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

class CrossAttentionFusion(nn.Module):
    """
    交叉注意力融合模块
    
    ConvNeXt特征（空间精细） ↔ CaiT tokens（全局语义）
    双向信息交换：全局上下文增强局部特征，局部细节引导全局注意力
    
    数学公式：
    Q_conv = W_q · F_conv   (ConvNet特征作为Query)
    K_trans = W_k · F_trans  (Transformer tokens作为Key)
    V_trans = W_v · F_trans  (Transformer tokens作为Value)
    
    F_cross = Softmax(Q_conv · K_trans^T / √d_k) · V_trans
    
    反向同理：Transformer用ConvNet特征做Query
    """
    
    def __init__(
        self, 
        conv_dim: int = 1024,    # ConvNeXt-Base 特征维度
        trans_dim: int = 384,   # CaiT-S36 特征维度
        num_heads: int = 8,
        dropout: float = 0.1
    ):
        super().__init__()
        self.num_heads = num_heads
        self.scale = (trans_dim // num_heads) ** -0.5
        
        # 投影到统一维度
        self.conv_proj = nn.Linear(conv_dim, trans_dim)
        self.trans_proj = nn.Linear(trans_dim, trans_dim)
        
        # 交叉注意力 QKV
        self.q_conv = nn.Linear(trans_dim, trans_dim)
        self.k_trans = nn.Linear(trans_dim, trans_dim)
        self.v_trans = nn.Linear(trans_dim, trans_dim)
        
        self.q_trans = nn.Linear(trans_dim, trans_dim)
        self.k_conv = nn.Linear(trans_dim, trans_dim)
        self.v_conv = nn.Linear(trans_dim, trans_dim)
        
        # 输出投影
        self.out_conv = nn.Linear(trans_dim, conv_dim)
        self.out_trans = nn.Linear(trans_dim, trans_dim)
        
        self.dropout = nn.Dropout(dropout)
        self.norm_conv = nn.LayerNorm(conv_dim)
        self.norm_trans = nn.LayerNorm(trans_dim)
        
    def forward(
        self, conv_feat: torch.Tensor, trans_tokens: torch.Tensor
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        Args:
            conv_feat: ConvNeXt特征 [B, N_conv, conv_dim]
            trans_tokens: CaiT tokens [B, N_trans, trans_dim]
            
        Returns:
            conv_enhanced: 增强后的Conv特征 [B, N_conv, conv_dim]
            trans_enhanced: 增强后的Trans tokens [B, N_trans, trans_dim]
        """
        B = conv_feat.shape[0]
        
        # 投影到统一维度
        conv_proj = self.conv_proj(conv_feat)  # [B, N_conv, trans_dim]
        trans_proj = self.trans_proj(trans_tokens)  # [B, N_trans, trans_dim]
        
        # --- 方向1: Conv特征用Trans做Key/Value ---
        q_c = self.q_conv(conv_proj)  # [B, N_conv, trans_dim]
        k_t = self.k_trans(trans_proj)  # [B, N_trans, trans_dim]
        v_t = self.v_trans(trans_proj)  # [B, N_trans, trans_dim]
        
        # 多头注意力
        q_c = q_c.view(B, -1, self.num_heads, trans_proj.size(-1) // self.num_heads).transpose(1, 2)
        k_t = k_t.view(B, -1, self.num_heads, trans_proj.size(-1) // self.num_heads).transpose(1, 2)
        v_t = v_t.view(B, -1, self.num_heads, trans_proj.size(-1) // self.num_heads).transpose(1, 2)
        
        attn_ct = (q_c @ k_t.transpose(-2, -1)) * self.scale
        attn_ct = attn_ct.softmax(dim=-1)
        out_ct = (attn_ct @ v_t).transpose(1, 2).reshape(B, -1, trans_proj.size(-1))
        
        # --- 方向2: Trans tokens用Conv做Key/Value ---
        q_t = self.q_trans(trans_proj)
        k_c = self.k_conv(conv_proj)
        v_c = self.v_conv(conv_proj)
        
        q_t = q_t.view(B, -1, self.num_heads, trans_proj.size(-1) // self.num_heads).transpose(1, 2)
        k_c = k_c.view(B, -1, self.num_heads, trans_proj.size(-1) // self.num_heads).transpose(1, 2)
        v_c = v_c.view(B, -1, self.num_heads, trans_proj.size(-1) // self.num_heads).transpose(1, 2)
        
        attn_tc = (q_t @ k_c.transpose(-2, -1)) * self.scale
        attn_tc = attn_tc.softmax(dim=-1)
        out_tc = (attn_tc @ v_c).transpose(1, 2).reshape(B, -1, trans_proj.size(-1))
        
        # 投影回各自维度 + 残差连接 + LayerNorm
        conv_enhanced = self.norm_conv(conv_feat + self.out_conv(self.dropout(out_ct)))
        trans_enhanced = self.norm_trans(trans_tokens + self.out_trans(self.dropout(out_tc)))
        
        return conv_enhanced, trans_enhanced


class SEBlock(nn.Module):
    """Squeeze-and-Excitation 通道重标定"""
    def __init__(self, channels: int, reduction: int = 16):
        super().__init__()
        self.squeeze = nn.AdaptiveAvgPool2d(1)
        self.excite = nn.Sequential(
            nn.Linear(channels, channels // reduction),
            nn.ReLU(inplace=True),
            nn.Linear(channels // reduction, channels),
            nn.Sigmoid()
        )
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        b, c, _, _ = x.shape
        s = self.squeeze(x).view(b, c)
        s = self.excite(s).view(b, c, 1, 1)
        return x * s


class CBAM(nn.Module):
    """Convolutional Block Attention Module (空间+通道注意力)"""
    def __init__(self, channels: int, reduction: int = 16, kernel_size: int = 7):
        super().__init__()
        self.channel_att = SEBlock(channels, reduction)
        self.spatial_att = nn.Sequential(
            nn.Conv2d(2, 1, kernel_size=kernel_size, padding=kernel_size // 2),
            nn.Sigmoid()
        )
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.channel_att(x)
        avg_out = torch.mean(x, dim=1, keepdim=True)
        max_out, _ = torch.max(x, dim=1, keepdim=True)
        s = torch.cat([avg_out, max_out], dim=1)
        s = self.spatial_att(s)
        return x * s


class FocalLoss(nn.Module):
    """Focal Loss - 解决疲劳数据类别不平衡"""
    def __init__(self, alpha: float = 0.25, gamma: float = 2.0):
        super().__init__()
        self.alpha = alpha
        self.gamma = gamma
    
    def forward(self, preds: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
        ce = F.cross_entropy(preds, targets, reduction='none')
        pt = torch.exp(-ce)
        loss = self.alpha * (1 - pt) ** self.gamma * ce
        return loss.mean()


class DrowSFormerXAI(nn.Module):
    """
    DrowSFormer-XAI 完整模型
    
    ConvNeXt-Base + CaiT-S36 + Cross-Attention + SE + CBAM + Mish
    参数量: ~159M, FLOPs: ~63.4G
    
    分类输出: 4类
    - 0: drowsy (疲劳)
    - 1: non-drowsy (清醒)
    - 2: closed eyes (闭眼)
    - 3: open eyes (睁眼)
    """
    
    def __init__(self, num_classes: int = 4):
        super().__init__()
        import timm
        
        # ConvNeXt-Base 分支 (局部空间特征)
        self.conv_branch = timm.create_model(
            'convnext_base', pretrained=True, num_classes=0
        )
        conv_dim = 1024  # ConvNeXt-Base 输出维度
        
        # CaiT-S36 分支 (全局语义特征)
        self.trans_branch = timm.create_model(
            'cait_s36', pretrained=True, num_classes=0
        )
        trans_dim = 384  # CaiT-S36 输出维度
        
        # 交叉注意力融合
        self.cross_attn = CrossAttentionFusion(
            conv_dim=conv_dim, trans_dim=trans_dim, num_heads=8
        )
        
        # 后处理模块
        self.se = SEBlock(conv_dim, reduction=16)
        self.cbam = CBAM(conv_dim, reduction=16)
        self.mish = nn.Mish()
        
        # 分类头
        self.classifier = nn.Sequential(
            nn.Linear(conv_dim + trans_dim, 512),
            nn.Mish(),
            nn.Dropout(0.3),
            nn.Linear(512, num_classes)
        )
        
        # 标签平滑
        self.label_smoothing = 0.1
        
    def forward(
        self, x: torch.Tensor, return_attention: bool = False
    ) -> torch.Tensor:
        """
        Args:
            x: 输入图像 [B, 3, 224, 224]
            return_attention: 是否返回Grad-CAM热力图
            
        Returns:
            logits: 分类logits [B, num_classes]
            (可选) heatmap: Grad-CAM热力图
        """
        # 双分支特征提取
        conv_feat = self.conv_branch(x)  # [B, 1024]
        trans_feat = self.trans_branch(x)  # [B, 384]
        
        # Reshape为序列形式
        conv_seq = conv_feat.unsqueeze(1)  # [B, 1, 1024]
        trans_seq = trans_feat.unsqueeze(1)  # [B, 1, 384]
        
        # 交叉注意力融合
        conv_enhanced, trans_enhanced = self.cross_attn(conv_seq, trans_seq)
        
        # Conv分支后处理
        conv_out = conv_enhanced.squeeze(1)  # [B, 1024]
        conv_2d = conv_out.view(-1, 1024, 1, 1)
        conv_2d = self.se(conv_2d)
        conv_2d = self.cbam(conv_2d)
        conv_2d = self.mish(conv_2d)
        conv_out = conv_2d.view(-1, 1024)
        
        trans_out = trans_enhanced.squeeze(1)  # [B, 384]
        
        # 拼接分类
        combined = torch.cat([conv_out, trans_out], dim=-1)  # [B, 1408]
        logits = self.classifier(combined)
        
        if return_attention:
            # Grad-CAM: 使用最后一个卷积层的梯度
            heatmap = self._compute_gradcam(conv_2d, logits)
            return logits, heatmap
        
        return logits
    
    def _compute_gradcam(
        self, feature_map: torch.Tensor, logits: torch.Tensor
    ) -> torch.Tensor:
        """Grad-CAM 可解释性热力图"""
        # 获取预测类别的梯度
        pred_class = logits.argmax(dim=1)
        one_hot = F.one_hot(pred_class, num_classes=logits.size(1)).float()
        
        grad = torch.autograd.grad(
            logits[:, pred_class].sum(),
            feature_map,
            retain_graph=True,
            create_graph=False
        )[0]
        
        # 全局平均池化梯度 → 通道权重
        weights = grad.mean(dim=(2, 3), keepdim=True)  # [B, C, 1, 1]
        
        # 加权求和
        heatmap = (weights * feature_map).sum(dim=1, keepdim=True)
        heatmap = F.relu(heatmap)
        heatmap = F.interpolate(
            heatmap, size=(224, 224), mode='bilinear', align_corners=False
        )
        heatmap = heatmap / (heatmap.max() + 1e-8)
        
        return heatmap


# ==================== 训练与评估 ====================
if __name__ == "__main__":
    # 模型配置
    model = DrowSFormerXAI(num_classes=4)
    
    # 统计参数量
    total_params = sum(p.numel() for p in model.parameters())
    trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
    print(f"总参数量: {total_params / 1e6:.1f}M")
    print(f"可训练参数: {trainable_params / 1e6:.1f}M")
    
    # 模拟输入
    x = torch.randn(4, 3, 224, 224)
    
    # 前向传播
    logits = model(x)
    print(f"输入: {x.shape}")
    print(f"输出: {logits.shape}")
    print(f"Logits: {logits[0]}")
    
    # 训练配置（按论文设置）
    optimizer = torch.optim.Adam(
        model.parameters(), 
        lr=1e-4,
        weight_decay=1e-4
    )
    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
        optimizer, T_max=50, eta_min=1e-6
    )
    criterion = FocalLoss(alpha=0.25, gamma=2.0)
    
    print(f"\n--- 训练配置 ---")
    print(f"优化器: Adam (lr=1e-4, weight_decay=1e-4)")
    print(f"调度器: CosineAnnealing (T_max=50)")
    print(f"损失函数: Focal Loss (alpha=0.25, gamma=2.0)")
    print(f"标签平滑: 0.1")
    print(f"Batch size: 64")
    print(f"最大epoch: 50 (early stop patience=10)")
    
    # 数据增强配置
    print(f"\n--- 数据增强 ---")
    print(f"旋转: ±30°")
    print(f"水平翻转: 50%")
    print(f"亮度: 50-150%")
    print(f"缩放: ±20%")
    print(f"平移: ±20%")

# 完整损失函数 = Focal Loss + Label Smoothing

class DrowSFormerLoss(nn.Module):
    """组合损失：Focal Loss + Label Smoothing"""
    
    def __init__(self, alpha=0.25, gamma=2.0, smoothing=0.1):
        super().__init__()
        self.focal = FocalLoss(alpha, gamma)
        self.smoothing = smoothing
        
    def forward(self, preds, targets):
        # Focal Loss
        focal_loss = self.focal(preds, targets)
        
        # Label Smoothing
        n_classes = preds.size(-1)
        log_probs = F.log_softmax(preds, dim=-1)
        smooth_loss = -log_probs.mean()
        
        return focal_loss + self.smoothing * smooth_loss

# 量化部署建议
# 原始模型: 159M参数, 63.4 GFLOPS
# 量化后: ~40M参数 (INT8), ~16 GFLOPS

# 步骤1: 知识蒸馏（CaiT-S36 → 轻量Transformer）
# 步骤2: INT8量化（FP32 → INT8）
# 步骤3: 剪枝（移除注意力权重小的head）
# 步骤4: ONNX导出 + TensorRT部署

# 目标芯片: QCS8255 (Hexagon NPU, 26 TOPS)
# 预期推理延迟: <15ms (量化后)
# 预期帧率: >60fps