import torch
import torch.nn as nn
from transformers import VideoMAEModel

class VideoEncoder(nn.Module):
    """
    EyeCue 视频编码器
    基于VideoMAE提取自视视频的时空特征
    """
    def __init__(self, hidden_dim: int = 768):
        super().__init__()
        # 使用预训练VideoMAE作为视频特征提取器
        self.video_mae = VideoMAEModel.from_pretrained("MCG-NJU/videomae-base")
        self.proj = nn.Linear(768, hidden_dim)
        
    def forward(self, video_frames: torch.Tensor) -> torch.Tensor:
        """
        Args:
            video_frames: (B, T, C, H, W) 
                         B=batch, T=时间步, C=通道, H/W=空间
        Returns:
            video_features: (B, T, hidden_dim)
        """
        B, T, C, H, W = video_frames.shape
        # 展平batch和时间维度
        x = video_frames.view(B * T, C, H, W)
        feats = self.video_mae(x).last_hidden_state  # (B*T, N_tokens, 768)
        # 取CLS token作为每帧特征
        feats = feats[:, 0, :]  # (B*T, 768)
        feats = self.proj(feats)  # (B*T, hidden_dim)
        return feats.view(B, T, -1)  # (B, T, hidden_dim)

class GazeEncoder(nn.Module):
    """
    EyeCue 注视编码器
    处理眼动时序数据，提取注视模式特征
    """
    def __init__(self, input_dim: int = 4, hidden_dim: int = 256, 
                 num_layers: int = 2, nhead: int = 4):
        super().__init__()
        # 输入: (x, y, pupil_diameter, timestamp_delta)
        self.input_proj = nn.Linear(input_dim, hidden_dim)
        
        # 使用Transformer编码器处理时序眼动数据
        encoder_layer = nn.TransformerEncoderLayer(
            d_model=hidden_dim,
            nhead=nhead,
            dim_feedforward=hidden_dim * 4,
            dropout=0.1,
            batch_first=True
        )
        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers)
        
        # 位置编码
        self.pos_encoding = PositionalEncoding(hidden_dim)
        
    def forward(self, gaze_seq: torch.Tensor, 
                mask: torch.Tensor = None) -> torch.Tensor:
        """
        Args:
            gaze_seq: (B, T, 4) 眼动序列
                      [x, y, pupil_diameter, dt]
            mask: (B, T) padding mask
        Returns:
            gaze_features: (B, T, hidden_dim)
        """
        x = self.input_proj(gaze_seq)  # (B, T, hidden_dim)
        x = self.pos_encoding(x)
        x = self.transformer(x, src_key_padding_mask=mask)
        return x

class GazeDrivenSemanticQuery(nn.Module):
    """
    注视驱动语义查询模块
    
    核心思想：根据注视点位置，从视频特征中
    动态选择驾驶员正在关注的视觉区域
    """
    def __init__(self, video_dim: int = 768, gaze_dim: int = 256,
                 num_queries: int = 8, nhead: int = 4):
        super().__init__()
        self.num_queries = num_queries
        
        # 将注视点映射到视觉token空间
        self.gaze_to_query = nn.Linear(gaze_dim, video_dim)
        
        # 多头交叉注意力
        self.cross_attn = nn.MultiheadAttention(
            embed_dim=video_dim,
            num_heads=nhead,
            kdim=video_dim,
            vdim=video_dim,
            batch_first=True
        )
        
        # 层归一化和FFN
        self.norm1 = nn.LayerNorm(video_dim)
        self.norm2 = nn.LayerNorm(video_dim)
        self.ffn = nn.Sequential(
            nn.Linear(video_dim, video_dim * 4),
            nn.GELU(),
            nn.Linear(video_dim * 4, video_dim)
        )
        
    def forward(self, video_feats: torch.Tensor, 
                gaze_feats: torch.Tensor) -> torch.Tensor:
        """
        Args:
            video_feats: (B, T_v, video_dim) 视频特征
            gaze_feats: (B, T_g, gaze_dim) 注视特征
        Returns:
            queried_features: (B, T_g, video_dim) 查询后的特征
        """
        # 用注视特征作为query查询视频特征
        queries = self.gaze_to_query(gaze_feats)  # (B, T_g, video_dim)
        
        # 交叉注意力：注视query -> 视频key/value
        attn_out, _ = self.cross_attn(
            query=queries,
            key=video_feats,
            value=video_feats
        )
        
        # 残差连接
        x = self.norm1(queries + attn_out)
        
        # FFN
        x = self.norm2(x + self.ffn(x))
        
        return x

class EyeCueModel(nn.Module):
    """
    EyeCue 完整模型
    
    论文：arXiv:2605.07859
    代码：https://github.com/langzhang2000/EyeCue
    """
    def __init__(self, config: dict):
        super().__init__()
        self.video_encoder = VideoEncoder(
            hidden_dim=config.get('video_dim', 768)
        )
        self.gaze_encoder = GazeEncoder(
            input_dim=config.get('gaze_input_dim', 4),
            hidden_dim=config.get('gaze_dim', 256),
            num_layers=config.get('gaze_layers', 2),
            nhead=config.get('nhead', 4)
        )
        self.gdsq = GazeDrivenSemanticQuery(
            video_dim=config.get('video_dim', 768),
            gaze_dim=config.get('gaze_dim', 256),
            num_queries=config.get('num_queries', 8),
            nhead=config.get('nhead', 4)
        )
        
        # 分类头
        fusion_dim = config.get('video_dim', 768) + config.get('gaze_dim', 256)
        self.classifier = nn.Sequential(
            nn.Linear(fusion_dim, 256),
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(256, 2)  # 认知分心 vs 正常
        )
        
    def forward(self, video_frames: torch.Tensor, 
                gaze_seq: torch.Tensor,
                gaze_mask: torch.Tensor = None) -> torch.Tensor:
        """
        Args:
            video_frames: (B, T, C, H, W)
            gaze_seq: (B, T_g, 4)
            gaze_mask: (B, T_g)
        Returns:
            logits: (B, 2)
        """
        # 1. 提取视频特征
        video_feats = self.video_encoder(video_frames)  # (B, T, 768)
        
        # 2. 提取注视特征
        gaze_feats = self.gaze_encoder(gaze_seq, gaze_mask)  # (B, T_g, 256)
        
        # 3. 注视驱动语义查询
        queried = self.gdsq(video_feats, gaze_feats)  # (B, T_g, 768)
        
        # 4. 融合: 拼接查询特征和原始注视特征
        # 对query特征做时序池化
        query_pooled = queried.mean(dim=1)  # (B, 768)
        gaze_pooled = gaze_feats.mean(dim=1)  # (B, 256)
        fused = torch.cat([query_pooled, gaze_pooled], dim=-1)  # (B, 1024)
        
        # 5. 分类
        logits = self.classifier(fused)
        return logits


# ==================== 测试代码 ====================
if __name__ == "__main__":
    import numpy as np
    
    config = {
        'video_dim': 768,
        'gaze_input_dim': 4,
        'gaze_dim': 256,
        'gaze_layers': 2,
        'nhead': 4,
        'num_queries': 8
    }
    
    model = EyeCueModel(config)
    
    # 模拟输入
    batch_size = 2
    video = torch.randn(batch_size, 16, 3, 224, 224)  # 16帧视频
    gaze = torch.randn(batch_size, 50, 4)  # 50个眼动采样点
    
    with torch.no_grad():
        logits = model(video, gaze)
        probs = torch.softmax(logits, dim=1)
        
    print(f"输出形状: {logits.shape}")
    print(f"分心概率: {probs[:, 1].tolist()}")
    print(f"参数量: {sum(p.numel() for p in model.parameters()) / 1e6:.1f}M")