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

class EyeCueModel(nn.Module):
    """
    EyeCue: 眼动+自我中心视频认知分心检测
    
    论文：IJCAI 2026, arXiv:2605.07859
    核心：跨模态交互建模眼动-场景关系
    
    输入：
        - video: (B, T, C, H, W) 自我中心视频
        - gaze: (B, T, 4) 眼动数据 (x, y, pupil_l, pupil_r)
    输出：
        - logits: (B, 2) 认知分心/正常
    """
    
    def __init__(self, hidden_dim=256, n_heads=4):
        super().__init__()
        
        # 1. 视觉编码器（轻量CNN）
        self.visual_encoder = 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((4, 4)),
        )
        self.visual_proj = nn.Linear(128 * 16, hidden_dim)
        
        # 2. 眼动编码器
        self.gaze_encoder = nn.Sequential(
            nn.Linear(4, 64), nn.ReLU(),
            nn.Linear(64, hidden_dim), nn.ReLU(),
        )
        
        # 3. 跨模态交互（Cross-Modal Attention）
        # 眼动作为Query，视觉作为Key/Value
        self.cross_attention = nn.MultiheadAttention(
            embed_dim=hidden_dim,
            num_heads=n_heads,
            batch_first=True,
        )
        
        # 4. 时序聚合
        self.temporal_lstm = nn.LSTM(
            input_size=hidden_dim * 2,  # 视觉+眼动交互后
            hidden_size=hidden_dim,
            num_layers=2,
            batch_first=True,
            dropout=0.1,
        )
        
        # 5. 分类头
        self.classifier = nn.Sequential(
            nn.Linear(hidden_dim, 64), nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(64, 2)
        )
    
    def forward(self, video, gaze):
        """
        Args:
            video: (B, T, C, H, W) 自我中心视频序列
            gaze: (B, T, 4) 眼动序列 (x, y, pupil_l, pupil_r)
        
        Returns:
            logits: (B, 2) 认知分心/正常
            attention_weights: 跨模态注意力权重（可解释性）
        """
        B, T, C, H, W = video.shape
        
        # 1. 视觉编码
        video_flat = video.view(B * T, C, H, W)
        vis_feat = self.visual_encoder(video_flat)  # (B*T, 128*16)
        vis_feat = self.visual_proj(vis_feat)  # (B*T, hidden)
        vis_feat = vis_feat.view(B, T, -1)  # (B, T, hidden)
        
        # 2. 眼动编码
        gaze_feat = self.gaze_encoder(gaze)  # (B, T, hidden)
        
        # 3. 跨模态交互（眼动Query → 视觉Key/Value）
        # 眼动作为Query询问视觉场景中哪个区域与当前注视相关
        attn_out, attn_weights = self.cross_attention(
            query=gaze_feat,    # (B, T, hidden)
            key=vis_feat,       # (B, T, hidden)
            value=vis_feat,     # (B, T, hidden)
        )
        
        # 拼接交互结果和原始视觉特征
        combined = torch.cat([attn_out, vis_feat], dim=-1)  # (B, T, hidden*2)
        
        # 4. 时序聚合
        lstm_out, _ = self.temporal_lstm(combined)  # (B, T, hidden)
        feat = lstm_out[:, -1, :]  # 取最后时刻 (B, hidden)
        
        # 5. 分类
        logits = self.classifier(feat)  # (B, 2)
        
        return {
            'logits': logits,
            'attention_weights': attn_weights,  # (B, T, T) 注意力热力图
        }


# 测试
if __name__ == "__main__":
    model = EyeCueModel()
    
    # 模拟输入：30秒@2fps = 60帧
    video = torch.randn(2, 60, 3, 64, 64)   # 60帧视频
    gaze = torch.randn(2, 60, 4)             # 60个眼动样本
    
    output = model(video, gaze)
    
    print(f"输入视频: {video.shape}")
    print(f"输入眼动: {gaze.shape}")
    print(f"分类输出: {output['logits'].shape}")
    print(f"注意力权重: {output['attention_weights'].shape}")
    print(f"参数总数: {sum(p.numel() for p in model.parameters()):,}")
    
    # 预测
    pred = F.softmax(output['logits'], dim=-1)
    print(f"\n预测概率: 正常={pred[0,0]:.2%}, 认知分心={pred[0,1]:.2%}")