"""
CNN-BiLSTM混合架构实现

CNN: 提取每帧的空间特征（面部关键点/姿态）
BiLSTM: 建模帧间时序依赖（行为序列）
注意力: 加权关键帧（如分心瞬间）

知识迁移：
1. 源域：轿车DMS数据（大量，如StateFarm/DDD数据集）
2. 目标域：电动公交车（少量）
3. 迁移策略：预训练CNN→微调BiLSTM
"""

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

class CNNBiLSTMDistraction(nn.Module):
    """
    CNN-BiLSTM分心驾驶检测框架
    
    输入: 视频帧序列 (B, T, C, H, W)
    输出: 分心类别概率 (B, n_classes)
    
    参数: 可配置CNN骨干和LSTM隐藏维度
    """
    
    def __init__(
        self,
        n_classes: int = 5,
        cnn_backbone: str = "mobilenet_v3_small",
        lstm_hidden: int = 256,
        lstm_layers: int = 2,
        dropout: float = 0.3,
    ):
        super().__init__()
        self.n_classes = n_classes
        
        # CNN骨干（预训练，可从轿车数据迁移）
        if cnn_backbone == "mobilenet_v3_small":
            import torchvision.models as models
            backbone = models.mobilenet_v3_small(weights=models.MobileNet_V3_Small_Weights.DEFAULT)
            self.cnn = nn.Sequential(*list(backbone.children())[:-1])
            cnn_out = 576  # MobileNetV3-Small特征维度
        else:
            # 简化CNN
            self.cnn = nn.Sequential(
                nn.Conv2d(3, 32, 3, stride=2, padding=1), nn.ReLU(),
                nn.Conv2d(32, 64, 3, stride=2, padding=1), nn.ReLU(),
                nn.Conv2d(64, 128, 3, stride=2, padding=1), nn.ReLU(),
                nn.AdaptiveAvgPool2d(1),
            )
            cnn_out = 128
        
        # BiLSTM时序建模
        self.lstm = nn.LSTM(
            input_size=cnn_out,
            hidden_size=lstm_hidden,
            num_layers=lstm_layers,
            batch_first=True,
            bidirectional=True,
            dropout=dropout if lstm_layers > 1 else 0,
        )
        
        # 注意力层
        self.attention = nn.Sequential(
            nn.Linear(lstm_hidden * 2, lstm_hidden),
            nn.Tanh(),
            nn.Linear(lstm_hidden, 1),
        )
        
        # 分类头
        self.classifier = nn.Sequential(
            nn.Linear(lstm_hidden * 2, 128),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(128, n_classes),
        )
    
    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        Args:
            x: (B, T, C, H, W) 视频帧序列
        Returns:
            logits: (B, n_classes) 分类结果
            attention_weights: (B, T) 注意力权重
        """
        B, T, C, H, W = x.shape
        
        # CNN提取每帧特征
        x = x.view(B * T, C, H, W)
        features = self.cnn(x)
        features = features.view(B, T, -1)  # (B, T, D)
        
        # BiLSTM时序建模
        lstm_out, _ = self.lstm(features)  # (B, T, 2*hidden)
        
        # 注意力加权
        attn_weights = self.attention(lstm_out)  # (B, T, 1)
        attn_weights = torch.softmax(attn_weights, dim=1)
        
        # 加权求和
        weighted = (lstm_out * attn_weights).sum(dim=1)  # (B, 2*hidden)
        
        # 分类
        logits = self.classifier(weighted)
        
        return logits, attn_weights.squeeze(-1)


class KnowledgeTransfer:
    """
    知识迁移策略
    
    源域: 轿车DMS数据集（大量）
    目标域: 电动公交车数据（少量）
    
    迁移方法: Fine-tuning with frozen CNN
    1. 在源域数据上预训练CNN-BiLSTM
    2. 冻结CNN，在目标域数据上微调BiLSTM
    3. 可选：渐进式解冻CNN最后几层
    """
    
    def __init__(self, model: CNNBiLSTMDistraction):
        self.model = model
    
    def stage1_pretrain(self, source_data, epochs=50):
        """Stage 1: 在轿车数据上预训练"""
        print("Stage 1: 在源域(轿车)数据上预训练")
        # 全参数训练
        for param in self.model.parameters():
            param.requires_grad = True
        # ... 训练循环 ...
        
    def stage2_finetune(self, target_data, epochs=20):
        """Stage 2: 在公交车数据上微调"""
        print("Stage 2: 在目标域(公交车)数据上微调")
        
        # 冻结CNN
        for param in self.model.cnn.parameters():
            param.requires_grad = False
        
        # 仅训练BiLSTM+注意力+分类头
        trainable = [
            p for p in self.model.parameters() if p.requires_grad
        ]
        print(f"  可训练参数: {sum(p.numel() for p in trainable):,}")
        print(f"  冻结参数: {sum(p.numel() for p in self.model.cnn.parameters()):,}")
        
    def stage3_progressive_unfreeze(self, target_data, epochs=10):
        """Stage 3: 渐进式解冻CNN最后层"""
        print("Stage 3: 渐进式解冻")
        # 解冻CNN最后1-2层
        for param in list(self.model.cnn.parameters())[-4:]:
            param.requires_grad = True


# ==================== 测试 ====================
if __name__ == "__main__":
    # 模型初始化
    model = CNNBiLSTMDistraction(
        n_classes=5,
        cnn_backbone="simple",  # 测试用简化CNN
        lstm_hidden=128,
        lstm_layers=2,
    )
    
    total_params = sum(p.numel() for p in model.parameters())
    print(f"=== CNN-BiLSTM分心检测模型 ===")
    print(f"总参数: {total_params:,}")
    
    # 模拟输入
    batch_size = 4
    seq_len = 30  # 1秒@30fps
    x = torch.randn(batch_size, seq_len, 3, 64, 64)
    
    # 前向
    logits, attn = model(x)
    print(f"\n输入: {x.shape} (B, T, C, H, W)")
    print(f"输出: {logits.shape} (B, n_classes)")
    print(f"注意力: {attn.shape} (B, T)")
    
    # 分心类别
    classes = ['正常驾驶', '手机使用', '交谈分心', '疲劳', '其他']
    probs = torch.softmax(logits, dim=-1)
    for i in range(batch_size):
        pred = classes[probs[i].argmax()]
        conf = probs[i].max().item()
        print(f"  样本{i}: {pred} ({conf:.1%})")
    
    # 知识迁移
    print(f"\n=== 知识迁移策略 ===")
    transfer = KnowledgeTransfer(model)
    transfer.stage1_pretrain(None)
    transfer.stage2_finetune(None)
    transfer.stage3_progressive_unfreeze(None)
    
    print(f"\n=== 商用车辆DMS特殊性 ===")
    print(f"  轿车DMS: 个体驾驶、多样场景、大量数据")
    print(f"  公交DMS: 职业驾驶员、固定路线、少量数据")
    print(f"  迁移价值: 轿车数据预训练→公交场景微调")
    print(f"  挑战: 摄像头位置不同、驾驶员坐姿不同、光照不同")

# 单帧vs序列
# 单帧CNN: 只看当前状态→容易误判
# CNN-BiLSTM: 看序列→区分"看手机"vs"低头拿东西"

# 时序窗口选择
window_options = {
    '0.5s (15帧)': '太短，无法区分短暂动作',
    '1.0s (30帧)': '平衡，推荐',
    '2.0s (60帧)': '更准但延迟增加',
    '5.0s (150帧)': '疲劳检测需要，但分心检测过长',
}