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

class DrowSFormerXAI(nn.Module):
    """
    DrowSFormer-XAI: 双分支交叉注意力疲劳检测网络
    
    架构:
    1. ConvNeXt-Base 分支: 局部空间特征 (眼睑/微表情)
    2. CaiT-S36 分支: 全局语义上下文 (头部姿态/光照)
    3. Cross-Attentive Fusion: 双向信息交换
    4. SE + CBAM: 通道+空间注意力增强
    5. 分类头: 4类 (drowsy/non-drowsy/closed-eye/open-eye)
    
    输入: (B, 3, 224, 224) 面部图像
    输出: (B, 4) 分类 logits + attention maps (可解释性)
    """
    
    def __init__(self, num_classes=4, drop_rate=0.1):
        super().__init__()
        
        # CNN 分支: ConvNeXt-Base (局部特征)
        self.conv_branch = ConvNeXtBase(
            in_chans=3,
            depths=[3, 3, 27, 3],
            dims=[128, 256, 512, 1024],
            drop_path_rate=0.1
        )
        
        # Transformer 分支: CaiT-S36 (全局上下文)
        self.transformer_branch = CaiTS36(
            img_size=224,
            patch_size=16,
            embed_dim=384,
            depth=36,
            num_heads=8,
            drop_path_rate=0.1
        )
        
        # 交叉注意力融合
        self.cross_attn = CrossAttentiveFusion(
            conv_dim=1024,    # ConvNeXt 输出维度
            trans_dim=384,    # CaiT 输出维度
            num_heads=8,
            hidden_dim=512
        )
        
        # SE 通道注意力
        self.se_block = SqueezeExcitation(512)
        
        # CBAM 空间注意力
        self.cbam = CBAM(512)
        
        # 分类头
        self.classifier = nn.Sequential(
            nn.LayerNorm(512),
            nn.Dropout(drop_rate),
            nn.Linear(512, 256),
            nn.GELU(),
            nn.Dropout(drop_rate),
            nn.Linear(256, num_classes)
        )
        
    def forward(self, x, return_attention=False):
        """
        Args:
            x: (B, 3, 224, 224) 面部图像
            return_attention: 是否返回注意力图
            
        Returns:
            logits: (B, 4) 分类
            attn_maps: 可选, 注意力热图
        """
        # 双分支特征提取
        conv_feat = self.conv_branch(x)        # (B, 1024, 7, 7)
        trans_feat = self.transformer_branch(x)  # (B, 384, 196)
        
        # 交叉注意力融合
        fused, attn_weights = self.cross_attn(conv_feat, trans_feat)
        
        # SE + CBAM 注意力增强
        fused = self.se_block(fused)
        fused = self.cbam(fused)
        
        # 全局平均池化
        pooled = F.adaptive_avg_pool2d(fused, 1).flatten(1)  # (B, 512)
        
        # 分类
        logits = self.classifier(pooled)
        
        if return_attention:
            return logits, attn_weights
        return logits


class ConvNeXtBase(nn.Module):
    """简化版 ConvNeXt-Base"""
    
    def __init__(self, in_chans=3, depths=[3,3,27,3], 
                 dims=[128,256,512,1024], drop_path_rate=0.1):
        super().__init__()
        
        self.stem = nn.Sequential(
            nn.Conv2d(in_chans, dims[0], 4, 4),
            nn.LayerNorm([dims[0], 56, 56])
        )
        
        self.stages = nn.ModuleList()
        for i in range(4):
            stage = nn.Sequential(*[
                ConvNeXtBlock(dims[i], 7, drop_path_rate)
                for _ in range(depths[i])
            ])
            self.stages.append(stage)
            
            if i < 3:
                self.stages.append(nn.Sequential(
                    nn.LayerNorm([dims[i], -1, -1]),
                    nn.Conv2d(dims[i], dims[i+1], 2, 2)
                ))
    
    def forward(self, x):
        x = self.stem(x)
        for stage in self.stages:
            x = stage(x)
        return x


class ConvNeXtBlock(nn.Module):
    """ConvNeXt 残差块"""
    
    def __init__(self, dim, kernel_size=7, drop_path=0.1):
        super().__init__()
        self.dwconv = nn.Conv2d(dim, dim, kernel_size, 
                                padding=kernel_size//2, groups=dim)
        self.norm = nn.LayerNorm([dim, -1, -1])
        self.pwconv1 = nn.Conv2d(dim, dim*4, 1)
        self.act = nn.GELU()
        self.pwconv2 = nn.Conv2d(dim*4, dim, 1)
        self.drop = nn.Dropout(drop_path)
        
    def forward(self, x):
        identity = x
        x = self.dwconv(x)
        x = self.norm(x)
        x = self.pwconv1(x)
        x = self.act(x)
        x = self.pwconv2(x)
        x = self.drop(x)
        return identity + x


class CaiTS36(nn.Module):
    """简化版 CaiT-S36"""
    
    def __init__(self, img_size=224, patch_size=16, embed_dim=384,
                 depth=36, num_heads=8, drop_path_rate=0.1):
        super().__init__()
        self.patch_embed = nn.Conv2d(3, embed_dim, patch_size, patch_size)
        num_patches = (img_size // patch_size) ** 2  # 196
        
        self.cls_token = nn.Parameter(torch.randn(1, 1, embed_dim))
        self.pos_embed = nn.Parameter(torch.randn(1, num_patches, embed_dim))
        
        self.blocks = nn.ModuleList([
            CaiTBlock(embed_dim, num_heads, drop_path_rate)
            for _ in range(depth)
        ])
        
        self.cls_attn_blocks = nn.ModuleList([
            ClassAttentionBlock(embed_dim, num_heads)
            for _ in range(2)
        ])
        
    def forward(self, x):
        B = x.shape[0]
        x = self.patch_embed(x)  # (B, D, 14, 14)
        x = x.flatten(2).transpose(1, 2)  # (B, 196, D)
        
        x = x + self.pos_embed
        
        for block in self.blocks:
            x = block(x)
        
        # Class attention
        cls = self.cls_token.expand(B, -1, -1)
        for block in self.cls_attn_blocks:
            cls = block(cls, x)
        
        return cls.transpose(1, 2).unsqueeze(-1)  # (B, D, 1, 1)


class CaiTBlock(nn.Module):
    def __init__(self, dim, num_heads, drop_path=0.1):
        super().__init__()
        self.norm1 = nn.LayerNorm(dim)
        self.attn = nn.MultiheadAttention(dim, num_heads, batch_first=True)
        self.norm2 = nn.LayerNorm(dim)
        self.ffn = nn.Sequential(
            nn.Linear(dim, dim*4),
            nn.GELU(),
            nn.Linear(dim*4, dim)
        )
        self.drop = nn.Dropout(drop_path)
        self.gamma = nn.Parameter(torch.ones(1) * 1e-5)
        
    def forward(self, x):
        normed = self.norm1(x)
        attn_out, _ = self.attn(normed, normed, normed)
        x = x + self.drop(self.gamma * attn_out)
        x = x + self.drop(self.gamma * self.ffn(self.norm2(x)))
        return x


class ClassAttentionBlock(nn.Module):
    def __init__(self, dim, num_heads):
        super().__init__()
        self.norm = nn.LayerNorm(dim)
        self.attn = nn.MultiheadAttention(dim, num_heads, batch_first=True)
        self.ffn = nn.Sequential(
            nn.Linear(dim, dim*4), nn.GELU(), nn.Linear(dim*4, dim)
        )
        
    def forward(self, cls, x):
        kv = torch.cat([cls, x], dim=1)
        normed = self.norm(kv)
        attn_out, _ = self.attn(normed[:, :1], normed, normed)
        cls = cls + attn_out
        cls = cls + self.ffn(self.norm(cls))
        return cls


class CrossAttentiveFusion(nn.Module):
    """
    核心模块: 双分支交叉注意力融合
    
    ConvNeXt 特征 → Query → CaiT 特征 (K,V) → 增强的 ConvNeXt 特征
    CaiT 特征 → Query → ConvNeXt 特征 (K,V) → 增强的 CaiT 特征
    """
    
    def __init__(self, conv_dim, trans_dim, num_heads=8, hidden_dim=512):
        super().__init__()
        
        # 维度对齐
        self.conv_proj = nn.Conv2d(conv_dim, hidden_dim, 1)
        self.trans_proj = nn.Linear(trans_dim, hidden_dim)
        
        # 双向交叉注意力
        self.conv2trans_attn = nn.MultiheadAttention(
            hidden_dim, num_heads, batch_first=True
        )
        self.trans2conv_attn = nn.MultiheadAttention(
            hidden_dim, num_heads, batch_first=True
        )
        
        # 融合投影
        self.fuse = nn.Sequential(
            nn.Conv2d(hidden_dim*2, hidden_dim, 1),
            nn.GELU(),
            nn.Conv2d(hidden_dim, hidden_dim, 3, padding=1),
            nn.GELU()
        )
        
    def forward(self, conv_feat, trans_feat):
        """
        Args:
            conv_feat: (B, C_conv, H, W) ConvNeXt 特征
            trans_feat: (B, C_trans, L) CaiT 特征序列
            
        Returns:
            fused: (B, hidden_dim, H, W) 融合特征
            attn_weights: 注意力权重 (可解释性)
        """
        B, C_c, H, W = conv_feat.shape
        L = trans_feat.shape[-1]
        
        # 维度对齐
        conv_proj = self.conv_proj(conv_feat)  # (B, D, H, W)
        trans_proj = self.trans_proj(trans_feat)  # (B, D, L)
        
        # Conv → Transformer 交叉注意力
        conv_tokens = conv_proj.flatten(2).transpose(1, 2)  # (B, H*W, D)
        trans_tokens = trans_proj.transpose(1, 2)  # (B, L, D)
        
        conv_enhanced, attn_c2t = self.conv2trans_attn(
            conv_tokens, trans_tokens, trans_tokens
        )
        
        # Transformer → Conv 交叉注意力
        trans_enhanced, attn_t2c = self.trans2conv_attn(
            trans_tokens, conv_tokens, conv_tokens
        )
        
        # 重塑并融合
        conv_enhanced = conv_enhanced.transpose(1, 2).reshape(B, -1, H, W)
        trans_enhanced = trans_enhanced.transpose(1, 2).unsqueeze(-1)
        
        # 双向融合
        fused = torch.cat([
            conv_enhanced,
            F.interpolate(trans_enhanced, size=(H, W), mode='bilinear')
        ], dim=1)
        
        fused = self.fuse(fused)
        
        attn_weights = {
            'conv2trans': attn_c2t,  # CNN 特征关注哪些 Transformer 特征
            'trans2conv': attn_t2c   # Transformer 特征关注哪些 CNN 特征
        }
        
        return fused, attn_weights


class SqueezeExcitation(nn.Module):
    """SE 通道注意力"""
    
    def __init__(self, channels, reduction=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):
        w = self.fc(x).unsqueeze(-1).unsqueeze(-1)
        return x * w


class CBAM(nn.Module):
    """CBAM: 通道+空间注意力"""
    
    def __init__(self, channels, reduction=16, kernel_size=7):
        super().__init__()
        # 通道注意力
        self.ca = nn.Sequential(
            nn.AdaptiveAvgPool2d(1),
            nn.Conv2d(channels, channels//reduction, 1),
            nn.ReLU(),
            nn.Conv2d(channels//reduction, channels, 1),
            nn.Sigmoid()
        )
        # 空间注意力
        self.sa = nn.Sequential(
            nn.Conv2d(2, 1, kernel_size, padding=kernel_size//2),
            nn.Sigmoid()
        )
        
    def forward(self, x):
        # 通道注意力
        ca = self.ca(x)
        x = x * ca
        # 空间注意力
        sa = self.sa(torch.cat([x.max(1)[0].unsqueeze(1),
                                x.mean(1).unsqueeze(1)], dim=1))
        x = x * sa
        return x


# 测试
if __name__ == "__main__":
    model = DrowSFormerXAI(num_classes=4)
    x = torch.randn(2, 3, 224, 224)
    
    logits, attn = model(x, return_attention=True)
    print(f"Input: {x.shape}")
    print(f"Output: {logits.shape}")  # (2, 4)
    print(f"Attention keys: {list(attn.keys())}")
    print(f"Params: {sum(p.numel() for p in model.parameters()):,}")

class DrowSFormerLoss(nn.Module):
    """
    DrowSFormer-XAI 联合损失
    
    Focal Loss: 解决类别不平衡
    Label Smoothing: 防止过拟合
    """
    
    def __init__(self, alpha=0.25, gamma=2.0, smoothing=0.1,
                 num_classes=4, class_weights=None):
        super().__init__()
        self.alpha = alpha
        self.gamma = gamma
        self.smoothing = smoothing
        self.num_classes = num_classes
        self.class_weights = class_weights
        
    def focal_loss(self, logits, targets):
        """Focal Loss: 聚焦难分类样本"""
        ce = F.cross_entropy(logits, targets, reduction='none',
                            weight=self.class_weights)
        pt = torch.exp(-ce)
        focal = self.alpha * (1 - pt) ** self.gamma * ce
        return focal.mean()
    
    def label_smoothing_loss(self, logits, targets):
        """Label Smoothing: 防止过度自信"""
        log_probs = F.log_softmax(logits, dim=-1)
        targets_onehot = F.one_hot(targets, self.num_classes).float()
        targets_smooth = (1 - self.smoothing) * targets_onehot + \
                         self.smoothing / self.num_classes
        return (-targets_smooth * log_probs).sum(dim=-1).mean()
    
    def forward(self, logits, targets):
        focal = self.focal_loss(logits, targets)
        smooth = self.label_smoothing_loss(logits, targets)
        return focal + 0.3 * smooth


if __name__ == "__main__":
    loss_fn = DrowSFormerLoss(alpha=0.25, gamma=2.0, smoothing=0.1)
    logits = torch.randn(4, 4)
    targets = torch.tensor([0, 1, 1, 2])
    loss = loss_fn(logits, targets)
    print(f"Loss: {loss.item():.4f}")

def drowsformer_deployment_optimization():
    """
    DrowSFormer-XAI 部署优化方案
    
    原始模型: 159M params, 63.4 GFLOPS
    目标: QCS8295 NPU, 30fps, < 5W
    """
    stages = {
        "Stage 1: 知识蒸馏": {
            "方法": "ConvNeXt-Tiny + CaiT-S12 → 学生模型",
            "参数量": "20-30M",
            "GFLOPS": "15-20",
            "预期精度": "97-98%",
            "时间": "1-2周"
        },
        "Stage 2: INT8 量化": {
            "方法": "Post-training quantization + QAT",
            "参数量": "20-30M (INT8)",
            "模型大小": "8-12 MB",
            "预期精度": "96-97%",
            "时间": "1周"
        },
        "Stage 3: NPU 部署": {
            "方法": "ONNX → SNPE/QNN 转换",
            "推理时间": "5-10 ms/帧",
            "帧率": "100-200 fps (远超30fps需求)",
            "功耗": "< 3W",
            "时间": "2-4周"
        },
        "Stage 4: 可解释性保留": {
            "方法": "保留 CBAM 空间注意力输出",
            "输出": "分类结果 + 注意力热图",
            "用途": "Euro NCAP 审计 + 误报分析",
            "时间": "内置"
        }
    }
    
    for stage, info in stages.items():
        print(f"\n{stage}:")
        for k, v in info.items():
            print(f"  {k}: {v}")
    
    return stages


if __name__ == "__main__":
    drowsformer_deployment_optimization()