import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import ViTModel

class CNNBiLSTMDistractionDetector(nn.Module):
    """
    CNN-BiLSTM 混合分心检测框架
    
    架构：
    1. CNN (MobileNetV3) 提取空间特征
    2. BiLSTM 建模时序动态
    3. 注意力机制聚焦关键帧
    4. 知识迁移从乘用车预训练模型
    
    适用场景：电动公交车驾驶员
    """
    
    def __init__(self, n_classes=10, hidden_dim=256, n_layers=2, 
                 pretrained=True):
        super().__init__()
        self.n_classes = n_classes
        self.hidden_dim = hidden_dim
        
        # 1. CNN特征提取 (MobileNetV3轻量化)
        self.cnn = MobileNetV3Encoder(pretrained=pretrained)
        feat_dim = self.cnn.output_dim  # 576
        
        # 2. 投影层（知识迁移适配）
        self.proj = nn.Sequential(
            nn.Linear(feat_dim, hidden_dim),
            nn.LayerNorm(hidden_dim),
            nn.GELU(),
            nn.Dropout(0.1)
        )
        
        # 3. BiLSTM时序建模
        self.bilstm = nn.LSTM(
            input_size=hidden_dim,
            hidden_size=hidden_dim,
            num_layers=n_layers,
            batch_first=True,
            bidirectional=True,
            dropout=0.2
        )
        
        # 4. 时序注意力
        self.temporal_attention = nn.Sequential(
            nn.Linear(hidden_dim * 2, hidden_dim),
            nn.Tanh(),
            nn.Linear(hidden_dim, 1)
        )
        
        # 5. 分类头
        self.classifier = nn.Sequential(
            nn.Linear(hidden_dim * 2, hidden_dim),
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(hidden_dim, n_classes)
        )
        
        # 知识迁移权重
        self.kd_alpha = nn.Parameter(torch.tensor(0.5))
        
    def forward(self, x):
        """
        Args:
            x: (B, T, 3, H, W) 视频帧序列
        
        Returns:
            logits: (B, n_classes)
            attention_weights: (B, T) 注意力权重
        """
        B, T, C, H, W = x.shape
        
        # 1. CNN逐帧提取
        x = x.view(B * T, C, H, W)
        cnn_feat = self.cnn(x)  # (B*T, feat_dim)
        cnn_feat = self.proj(cnn_feat)  # (B*T, hidden_dim)
        cnn_feat = cnn_feat.view(B, T, -1)  # (B, T, hidden_dim)
        
        # 2. BiLSTM时序建模
        lstm_out, _ = self.bilstm(cnn_feat)  # (B, T, 2*hidden_dim)
        
        # 3. 时序注意力
        attn_scores = self.temporal_attention(lstm_out)  # (B, T, 1)
        attn_weights = F.softmax(attn_scores, dim=1)  # (B, T, 1)
        
        # 加权求和
        weighted = (lstm_out * attn_weights).sum(dim=1)  # (B, 2*hidden_dim)
        
        # 4. 分类
        logits = self.classifier(weighted)
        
        return logits, attn_weights.squeeze(-1)
    
    def knowledge_distillation_loss(self, student_logits, 
                                     teacher_logits, labels, 
                                     temperature=4.0):
        """知识蒸馏损失"""
        alpha = torch.sigmoid(self.kd_alpha)
        
        # 硬标签损失
        hard_loss = F.cross_entropy(student_logits, labels)
        
        # 软标签损失 (KL散度)
        soft_student = F.log_softmax(student_logits / temperature, dim=1)
        soft_teacher = F.softmax(teacher_logits / temperature, dim=1)
        soft_loss = F.kl_div(soft_student, soft_teacher, 
                             reduction='batchmean') * (temperature ** 2)
        
        # 总损失
        total_loss = alpha * hard_loss + (1 - alpha) * soft_loss
        return total_loss, {'hard': hard_loss.item(), 
                           'soft': soft_loss.item(),
                           'alpha': alpha.item()}


class MobileNetV3Encoder(nn.Module):
    """轻量化CNN编码器"""
    
    def __init__(self, pretrained=True):
        super().__init__()
        # 简化版MobileNetV3
        self.features = nn.Sequential(
            # Stem
            nn.Conv2d(3, 16, 3, stride=2, padding=1),
            nn.BatchNorm2d(16),
            nn.Hardswish(),
            
            # Blocks
            self._make_block(16, 16, 3, 1),
            self._make_block(16, 24, 3, 2),
            self._make_block(24, 24, 3, 1),
            self._make_block(24, 40, 5, 2),
            self._make_block(40, 40, 5, 1),
            self._make_block(40, 80, 3, 2),
            self._make_block(80, 80, 3, 1),
            self._make_block(80, 112, 3, 1),
            self._make_block(112, 160, 5, 2),
            self._make_block(160, 160, 5, 1),
            
            # Final
            nn.Conv2d(160, 576, 1),
            nn.BatchNorm2d(576),
            nn.Hardswish(),
            nn.AdaptiveAvgPool2d(1),
        )
        self.output_dim = 576
        
    def _make_block(self, in_ch, out_ch, kernel_size, stride):
        return nn.Sequential(
            nn.Conv2d(in_ch, out_ch, kernel_size, stride=stride, 
                     padding=kernel_size//2, groups=in_ch),
            nn.BatchNorm2d(out_ch),
            nn.Hardswish(),
            nn.Conv2d(out_ch, out_ch, 1),
            nn.BatchNorm2d(out_ch),
            nn.Hardswish(),
        )
    
    def forward(self, x):
        return self.features(x).flatten(1)


# 分心行为类别（公交车特定）
BUS_DISTRACTION_CLASSES = {
    0: 'safe_driving',        # 安全驾驶
    1: 'phone_call_right',    # 右手打电话
    2: 'phone_call_left',     # 左手打电话
    3: 'texting_right',       # 右手发短信
    4: 'texting_left',        # 左手发短信
    5: 'adjusting_radio',     # 调收音机/广播
    6: 'drinking',            # 喝水
    7: 'talking_passenger',   # 与乘客交谈
    8: 'operating_tickets',   # 操作票务系统
    9: 'reaching_back',       # 向后取物
}

# 测试
if __name__ == "__main__":
    B, T = 4, 16  # 4个样本, 16帧
    model = CNNBiLSTMDistractionDetector(n_classes=10)
    
    # 模拟视频输入
    video = torch.randn(B, T, 3, 224, 224)
    
    logits, attn = model(video)
    
    print(f"输入视频: {video.shape}")
    print(f"输出logits: {logits.shape}")
    print(f"注意力权重: {attn.shape}")
    print(f"参数量: {sum(p.numel() for p in model.parameters())/1e6:.2f}M")
    
    # 模拟知识蒸馏
    teacher_logits = torch.randn(B, 10)
    labels = torch.randint(0, 10, (B,))
    loss, info = model.knowledge_distillation_loss(logits, teacher_logits, labels)
    print(f"\n知识蒸馏损失: {loss.item():.4f}")
    print(f"  硬标签: {info['hard']:.4f}")
    print(f"  软标签: {info['soft']:.4f}")
    print(f"  alpha: {info['alpha']:.4f}")

KNOWLEDGE_TRANSFER_STRATEGY = {
    'stage1': {
        'action': '预训练',
        'source': '乘用车DMS数据集 (如State Farm Distracted Driver)',
        'target': 'CNN特征提取器',
        'method': '标准分类训练',
    },
    'stage2': {
        'action': '微调',
        'source': '公交车标注数据 (少量, ~1000样本)',
        'target': '全网络微调 + 知识蒸馏',
        'method': '教师(乘用车) → 学生(公交车)',
    },
    'stage3': {
        'action': '部署',
        'source': '微调后模型',
        'target': '边缘设备 (QCS8255)',
        'method': 'INT8量化 + TensorRT',
    }
}