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

class GDSQModule(nn.Module):
    """
    Gaze-Driven Semantic Query (GDSQ) 模块
    
    核心思想：利用凝视信息动态选择视觉 Token，
    反映驾驶员注意力在上下文中的分配模式
    
    论文 Section 3.3 实现
    """
    
    def __init__(
        self,
        visual_dim: int = 768,
        gaze_dim: int = 256,
        num_heads: int = 8,
        dropout: float = 0.1
    ):
        super().__init__()
        self.visual_dim = visual_dim
        self.gaze_dim = gaze_dim
        self.num_heads = num_heads
        
        # 凝视特征投影层
        self.gaze_projection = nn.Linear(gaze_dim, visual_dim)
        
        # 交叉注意力层（凝视 query -> 视觉 key/value）
        self.cross_attention = nn.MultiheadAttention(
            embed_dim=visual_dim,
            num_heads=num_heads,
            dropout=dropout,
            batch_first=True
        )
        
        # 前馈网络
        self.ffn = nn.Sequential(
            nn.Linear(visual_dim, visual_dim * 4),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(visual_dim * 4, visual_dim),
            nn.Dropout(dropout)
        )
        
        # Layer Norm
        self.norm1 = nn.LayerNorm(visual_dim)
        self.norm2 = nn.LayerNorm(visual_dim)
        
    def forward(
        self,
        visual_tokens: torch.Tensor,  # (B, T, L, D) - 批次, 时间, 空间位置, 维度
        gaze_features: torch.Tensor  # (B, T, D_g) - 批次, 时间, 凝视维度
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        前向传播
        
        Args:
            visual_tokens: 视频编码器输出的视觉 Token 序列
            gaze_features: 凝视编码器输出的凝视特征
            
        Returns:
            gaze_selected_features: 凝视选择的视觉特征
            global_features: 全局视觉特征
        """
        B, T, L, D = visual_tokens.shape
        
        # 将凝视特征投影到视觉维度
        gaze_query = self.gaze_projection(gaze_features)  # (B, T, D)
        gaze_query = gaze_query.unsqueeze(2)  # (B, T, 1, D)
        
        # 展平空间维度用于注意力
        visual_flat = visual_tokens.reshape(B * T, L, D)
        gaze_flat = gaze_query.reshape(B * T, 1, D)
        
        # 交叉注意力：凝视作为 query, 视觉 Token 作为 key/value
        attended, attn_weights = self.cross_attention(
            query=gaze_flat,
            key=visual_flat,
            value=visual_flat
        )
        
        # 残差连接 + Layer Norm
        attended = self.norm1(attended + gaze_flat)
        
        # FFN
        attended = self.norm2(self.ffn(attended) + attended)
        
        # 重塑回原始维度
        gaze_selected = attended.reshape(B, T, D)
        
        # 全局视觉特征（平均池化）
        global_features = visual_tokens.mean(dim=2)  # (B, T, D)
        
        return gaze_selected, global_features

import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import Tuple, Optional
import numpy as np

class VideoEncoder(nn.Module):
    """
    视频编码器 - 提取第一人称驾驶视频的上下文特征
    
    基于 Vision Transformer 的时序扩展
    """
    
    def __init__(
        self,
        frame_size: int = 224,
        patch_size: int = 16,
        in_channels: int = 3,
        embed_dim: int = 768,
        num_frames: int = 16,
        num_heads: int = 12,
        num_layers: int = 6
    ):
        super().__init__()
        self.num_frames = num_frames
        self.embed_dim = embed_dim
        
        # Patch Embedding
        num_patches = (frame_size // patch_size) ** 2
        self.patch_embed = nn.Conv2d(
            in_channels, embed_dim,
            kernel_size=patch_size, stride=patch_size
        )
        
        # 位置编码
        self.pos_embed = nn.Parameter(
            torch.randn(1, num_patches + 1, embed_dim) * 0.02
        )
        self.cls_token = nn.Parameter(torch.randn(1, 1, embed_dim) * 0.02)
        
        # 空间 Transformer
        encoder_layer = nn.TransformerEncoderLayer(
            d_model=embed_dim,
            nhead=num_heads,
            dim_feedforward=embed_dim * 4,
            dropout=0.1,
            batch_first=True
        )
        self.spatial_encoder = nn.TransformerEncoder(encoder_layer, num_layers)
        
        # 时序聚合
        self.temporal_pool = nn.AdaptiveAvgPool1d(1)
        
    def forward(self, video: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        Args:
            video: (B, C, T, H, W) - 批次, 通道, 帧数, 高, 宽
            
        Returns:
            visual_tokens: (B, T, L, D) - 每帧的视觉Token序列
            global_feature: (B, T, D) - 每帧的全局表示
        """
        B, C, T, H, W = video.shape
        
        # 逐帧提取 Patch 特征
        frames = video.permute(0, 2, 1, 3, 4)  # (B, T, C, H, W)
        frames = frames.reshape(B * T, C, H, W)
        
        patches = self.patch_embed(frames)  # (B*T, D, H/P, W/P)
        patches = patches.flatten(2).transpose(1, 2)  # (B*T, L, D)
        
        # 添加 CLS Token
        cls = self.cls_token.expand(B * T, -1, -1)
        patches = torch.cat([cls, patches], dim=1)
        patches = patches + self.pos_embed
        
        # 空间编码
        encoded = self.spatial_encoder(patches)
        
        # 重塑回时序维度
        L = encoded.shape[1]
        visual_tokens = encoded.reshape(B, T, L, -1)
        
        # 全局特征（CLS Token）
        global_feature = visual_tokens[:, :, 0, :]  # (B, T, D)
        
        return visual_tokens, global_feature


class GazeEncoder(nn.Module):
    """
    凝视编码器 - 提取眼动行为时序模式
    
    输入：注视点坐标序列 + 瞳孔信息
    输出：凝视行为特征向量
    """
    
    def __init__(
        self,
        gaze_input_dim: int = 4,  # x, y, pupil_diameter, validity
        hidden_dim: int = 256,
        output_dim: int = 768,
        num_layers: int = 2,
        num_heads: int = 4
    ):
        super().__init__()
        
        # 输入投影
        self.input_proj = nn.Linear(gaze_input_dim, hidden_dim)
        
        # 位置编码
        self.pos_encoding = nn.Parameter(
            torch.randn(1, 512, hidden_dim) * 0.02
        )
        
        # Transformer 编码器
        encoder_layer = nn.TransformerEncoderLayer(
            d_model=hidden_dim,
            nhead=num_heads,
            dim_feedforward=hidden_dim * 4,
            dropout=0.1,
            batch_first=True
        )
        self.encoder = nn.TransformerEncoder(encoder_layer, num_layers)
        
        # 输出投影到视觉维度
        self.output_proj = nn.Linear(hidden_dim, output_dim)
        
    def forward(self, gaze_seq: torch.Tensor) -> torch.Tensor:
        """
        Args:
            gaze_seq: (B, T, gaze_input_dim) - 注视点序列
            
        Returns:
            gaze_features: (B, T, D) - 凝视行为特征
        """
        x = self.input_proj(gaze_seq)
        x = x + self.pos_encoding[:, :x.shape[1], :]
        x = self.encoder(x)
        return self.output_proj(x)


class EyeCueModel(nn.Module):
    """
    EyeCue 完整模型
    
    论文：Driver Cognitive Distraction Detection via 
          Gaze-Empowered Egocentric Video Understanding
    会议：IJCAI 2026
    
    架构：视频编码器 + 凝视编码器 + GDSQ模块 + 融合分类器
    """
    
    def __init__(
        self,
        frame_size: int = 224,
        embed_dim: int = 768,
        num_frames: int = 16,
        gaze_dim: int = 4,
        num_classes: int = 2,
        dropout: float = 0.1
    ):
        super().__init__()
        
        # 视频编码器
        self.video_encoder = VideoEncoder(
            frame_size=frame_size,
            embed_dim=embed_dim,
            num_frames=num_frames
        )
        
        # 凝视编码器
        self.gaze_encoder = GazeEncoder(
            gaze_input_dim=gaze_dim,
            output_dim=embed_dim
        )
        
        # GDSQ 模块
        self.gdsq = GDSQModule(
            visual_dim=embed_dim,
            gaze_dim=embed_dim,
            num_heads=8,
            dropout=dropout
        )
        
        # 时序聚合
        self.temporal_attn = nn.MultiheadAttention(
            embed_dim=embed_dim * 3,  # gaze_selected + global + gaze
            num_heads=8,
            dropout=dropout,
            batch_first=True
        )
        
        # 分类头
        self.classifier = nn.Sequential(
            nn.Linear(embed_dim * 3, embed_dim),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(embed_dim, num_classes)
        )
        
    def forward(
        self,
        video: torch.Tensor,
        gaze_seq: torch.Tensor
    ) -> torch.Tensor:
        """
        Args:
            video: (B, C, T, H, W) - 第一人称驾驶视频
            gaze_seq: (B, T, 4) - 眼动追踪序列 (x, y, pupil, validity)
            
        Returns:
            logits: (B, num_classes) - 分心/正常分类
        """
        # 1. 视频编码
        visual_tokens, global_feat = self.video_encoder(video)
        
        # 2. 凝视编码
        gaze_feat = self.gaze_encoder(gaze_seq)
        
        # 3. GDSQ：凝视驱动的视觉 Token 选择
        gaze_selected, _ = self.gdsq(visual_tokens, gaze_feat)
        
        # 4. 特征拼接
        fused = torch.cat([gaze_selected, global_feat, gaze_feat], dim=-1)
        
        # 5. 时序注意力聚合
        fused_t = fused.permute(1, 0, 2)  # (T, B, D)
        pooled, _ = self.temporal_attn(fused_t, fused_t, fused_t)
        pooled = pooled.mean(dim=0)  # (B, D*3)
        
        # 6. 分类
        logits = self.classifier(pooled)
        
        return logits


# ============ 训练与测试 ============

def train_step(
    model: EyeCueModel,
    video: torch.Tensor,
    gaze: torch.Tensor,
    labels: torch.Tensor,
    optimizer: torch.optim.Optimizer,
    criterion: nn.Module,
    device: str = 'cuda'
) -> float:
    """单步训练"""
    model.train()
    video = video.to(device)
    gaze = gaze.to(device)
    labels = labels.to(device)
    
    optimizer.zero_grad()
    logits = model(video, gaze)
    loss = criterion(logits, labels)
    loss.backward()
    optimizer.step()
    
    return loss.item()


def evaluate(
    model: EyeCueModel,
    dataloader,
    device: str = 'cuda'
) -> dict:
    """评估模型"""
    model.eval()
    all_preds = []
    all_labels = []
    
    with torch.no_grad():
        for batch in dataloader:
            video, gaze, labels = batch
            video = video.to(device)
            gaze = gaze.to(device)
            
            logits = model(video, gaze)
            preds = torch.argmax(logits, dim=1)
            
            all_preds.extend(preds.cpu().numpy())
            all_labels.extend(labels.numpy())
    
    all_preds = np.array(all_preds)
    all_labels = np.array(all_labels)
    
    accuracy = (all_preds == all_labels).mean()
    
    # 计算每个类别的 precision/recall
    results = {'accuracy': accuracy}
    for cls in [0, 1]:
        tp = ((all_preds == cls) & (all_labels == cls)).sum()
        fp = ((all_preds == cls) & (all_labels != cls)).sum()
        fn = ((all_preds != cls) & (all_labels == cls)).sum()
        precision = tp / (tp + fp + 1e-8)
        recall = tp / (tp + fn + 1e-8)
        f1 = 2 * precision * recall / (precision + recall + 1e-8)
        results[f'class_{cls}_precision'] = precision
        results[f'class_{cls}_recall'] = recall
        results[f'class_{cls}_f1'] = f1
    
    return results


if __name__ == "__main__":
    # 模型初始化
    model = EyeCueModel(
        frame_size=224,
        embed_dim=384,  # 轻量版
        num_frames=16,
        gaze_dim=4,
        num_classes=2
    )
    
    # 参数量统计
    total_params = sum(p.numel() for p in model.parameters())
    print(f"EyeCue 模型参数量: {total_params:,}")
    print(f"模型大小: {total_params * 4 / 1024 / 1024:.1f} MB (FP32)")
    
    # 模拟输入测试
    B = 4  # 批次大小
    video = torch.randn(B, 3, 16, 224, 224)  # 16帧 224x224 视频
    gaze = torch.randn(B, 16, 4)  # 16帧凝视数据
    
    # 前向传播
    with torch.no_grad():
        logits = model(video, gaze)
    
    print(f"\n输入视频形状: {video.shape}")
    print(f"输入凝视形状: {gaze.shape}")
    print(f"输出 logits 形状: {logits.shape}")
    print(f"预测类别: {torch.argmax(logits, dim=1).tolist()}")
    
    # 训练配置
    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4, weight_decay=0.01)
    criterion = nn.CrossEntropyLoss()
    
    # 模拟训练步
    labels = torch.randint(0, 2, (B,))
    loss = train_step(model, video, gaze, labels, optimizer, criterion, device='cpu')
    print(f"\n训练损失: {loss:.4f}")
    
    print("\n=== 论文报告结果 ===")
    print("EyeCue 在 CogDrive 上的准确率: 74.38%")
    print("超越 11 个基线模型 > 7%")
    print("各场景准确率均 > 70%")

from torch.utils.data import Dataset, DataLoader
import json
from pathlib import Path

class CogDriveDataset(Dataset):
    """
    CogDrive 数据集加载器
    
    数据集结构:
    CogDrive/
    ├── DR(eye)VE/
    │   ├── video/
    │   └── gaze/
    ├── BDD-A/
    ├── DADA-2000/
    └── TrafficGaze/
    
    每个样本包含:
    - egocentric_video: 第一人称驾驶视频片段
    - gaze_sequence: 眼动追踪序列
    - label: 0=正常, 1=认知分心
    - scenario: 道路类型/天气/时间
    """
    
    def __init__(
        self,
        root: str = "./CogDrive",
        split: str = "train",
        num_frames: int = 16,
        frame_size: int = 224,
        augment: bool = True
    ):
        self.root = Path(root)
        self.split = split
        self.num_frames = num_frames
        self.frame_size = frame_size
        self.augment = augment and split == "train"
        
        # 加载标注文件
        ann_file = self.root / f"{split}_annotations.json"
        with open(ann_file) as f:
            self.annotations = json.load(f)
    
    def __len__(self):
        return len(self.annotations)
    
    def __getitem__(self, idx):
        ann = self.annotations[idx]
        
        # 加载视频帧
        video = self._load_video(ann['video_path'])
        
        # 加载凝视数据
        gaze = self._load_gaze(ann['gaze_path'])
        
        # 标签
        label = ann['label']  # 0=正常, 1=认知分心
        
        return video, gaze, label
    
    def _load_video(self, path: str) -> torch.Tensor:
        """加载视频并预处理"""
        # 实际实现需用 cv2/decord 加载视频
        # 这里给出预处理流程
        import cv2
        cap = cv2.VideoCapture(str(self.root / path))
        frames = []
        while len(frames) < self.num_frames:
            ret, frame = cap.read()
            if not ret:
                break
            frame = cv2.resize(frame, (self.frame_size, self.frame_size))
            frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            frames.append(frame)
        cap.release()
        
        # 均匀采样到 num_frames
        if len(frames) > self.num_frames:
            indices = np.linspace(0, len(frames)-1, self.num_frames, dtype=int)
            frames = [frames[i] for i in indices]
        elif len(frames) < self.num_frames:
            # 重复最后一帧
            frames += [frames[-1]] * (self.num_frames - len(frames))
        
        # 归一化 + 增强
        video = torch.stack([
            torch.from_numpy(f).float().permute(2, 0, 1) / 255.0
            for f in frames
        ])  # (T, C, H, W)
        video = video.permute(1, 0, 2, 3)  # (C, T, H, W)
        
        # 标准化
        mean = torch.tensor([0.485, 0.456, 0.406]).view(3, 1, 1, 1)
        std = torch.tensor([0.229, 0.224, 0.225]).view(3, 1, 1, 1)
        video = (video - mean) / std
        
        return video
    
    def _load_gaze(self, path: str) -> torch.Tensor:
        """加载凝视数据"""
        # 格式: [x, y, pupil_diameter, validity] x T
        gaze = np.load(str(self.root / path))
        
        # 归一化坐标到 [0, 1]
        gaze[:, 0] /= 1920  # 图像宽度
        gaze[:, 1] /= 1080  # 图像高度
        
        # 对齐到视频帧数
        if len(gaze) > self.num_frames:
            indices = np.linspace(0, len(gaze)-1, self.num_frames, dtype=int)
            gaze = gaze[indices]
        elif len(gaze) < self.num_frames:
            gaze = np.pad(gaze, ((0, self.num_frames - len(gaze)), (0, 0)), 
                         mode='edge')
        
        return torch.from_numpy(gaze).float()

class CognitiveDistractionDetector:
    """
    IMS 认知分心检测模块
    
    基于 EyeCue 思路简化，适配车内 DMS 摄像头场景
    """
    
    def __init__(self, config: dict):
        # 简化版：使用车内摄像头而非第一人称视频
        # + 眼动追踪输出
        self.gaze_history = []
        self.scene_context = None
        self.window_size = config.get('window_sec', 10)
        self.fps = config.get('fps', 30)
        self.threshold = config.get('threshold', 0.5)
        
    def update(
        self,
        gaze_x: float,
        gaze_y: float,
        scene_descriptor: np.ndarray,
        driving_state: str
    ) -> tuple:
        """
        每帧更新
        
        Args:
            gaze_x, gaze_y: 注视点坐标 (归一化)
            scene_descriptor: 场景特征向量 (来自前视摄像头)
            driving_state: 驾驶状态 ("highway"/"urban"/"parked")
            
        Returns:
            (is_distracted, confidence)
        """
        self.gaze_history.append((gaze_x, gaze_y, scene_descriptor, driving_state))
        
        if len(self.gaze_history) > self.window_size * self.fps:
            self.gaze_history.pop(0)
        
        if len(self.gaze_history) < self.window_size * self.fps:
            return False, 0.0
        
        # 计算凝视-上下文一致性
        history = np.array([(g[0], g[1]) for g in self.gaze_history])
        
        # 特征1：注视点方差（认知分心时增大）
        gaze_var = np.var(history, axis=0).sum()
        
        # 特征2：注视点与道路中心的偏离
        road_center = np.array([0.5, 0.6])  # 假设道路中心
        gaze_deviation = np.mean(np.linalg.norm(history - road_center, axis=1))
        
        # 特征3：扫视频率（认知分心时模式异常）
        saccades = np.diff(history, axis=0)
        saccade_freq = np.mean(np.linalg.norm(saccades, axis=1))
        
        # 综合评分
        distraction_score = (
            0.4 * min(gaze_var * 10, 1.0) +
            0.3 * min(gaze_deviation * 2, 1.0) +
            0.3 * min(saccade_freq * 5, 1.0)
        )
        
        is_distracted = distraction_score > self.threshold
        return is_distracted, distraction_score