# models/video_encoder.py 核心结构
import torch
import torch.nn as nn
from transformers import TimesformerModel

class VideoEncoder(nn.Module):
    """
    TimeSformer 骨干网络，支持任意 clip 长度
    通过时间嵌入插值适配不同帧数
    """
    def __init__(self, clip_len=8):
        super().__init__()
        self.timesformer = TimesformerModel.from_pretrained(
            "facebook/timesformer-base-finetuned-k600"
        )
        self.clip_len = clip_len
        # 原始模型支持 16 帧，通过插值扩展
        
    def forward(self, video_frames):
        """
        Args:
            video_frames: (B, C, T, H, W) 视频帧序列
        Returns:
            cls_token: (B, 1, D) 全局表示
            video_tokens: (B, T*196, D) 空间 token
            # 196 = 14*14 patches, 每帧的空间 patch 数
        """
        outputs = self.timesformer(video_frames)
        cls_token = outputs.last_hidden_state[:, :1, :]
        video_tokens = outputs.last_hidden_state[:, 1:, :]
        return cls_token, video_tokens

# models/gaze_encoder.py 核心结构
class GazeEncoder(nn.Module):
    """
    线性投影 + 可学习 CLS Token + 自注意力
    输入: 原始凝视坐标序列 (B, T, 2)
    输出: 凝视表示 (B, T, D) 和 CLS 表示 (B, 1, D)
    """
    def __init__(self, embed_dim=768, num_heads=12, num_layers=4):
        super().__init__()
        # 线性投影: 2D 坐标 → 高维空间
        self.proj = nn.Linear(2, embed_dim)
        # 可学习 CLS Token
        self.cls_token = nn.Parameter(torch.randn(1, 1, embed_dim))
        # 多层自注意力
        encoder_layer = nn.TransformerEncoderLayer(
            d_model=embed_dim,
            nhead=num_heads,
            dim_feedforward=3072,
            dropout=0.1,
            batch_first=True
        )
        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers)
        
    def forward(self, gaze_sequence):
        """
        Args:
            gaze_sequence: (B, T, 2) 凝视点坐标
        Returns:
            cls_out: (B, 1, D) CLS 表示
            gaze_tokens: (B, T, D) 时序凝视表示
        """
        B = gaze_sequence.shape[0]
        x = self.proj(gaze_sequence)  # (B, T, D)
        
        # 拼接 CLS Token
        cls = self.cls_token.expand(B, -1, -1)  # (B, 1, D)
        x = torch.cat([cls, x], dim=1)  # (B, T+1, D)
        
        x = self.transformer(x)
        cls_out = x[:, :1, :]
        gaze_tokens = x[:, 1:, :]
        return cls_out, gaze_tokens

# models/semantic.py 核心逻辑
class GazeGuidedPatchSelection:
    """
    将凝视坐标映射到 TimeSformer 的 patch 空间
    
    视频帧被划分为 14×14 = 196 个 patch
    凝视点 (x, y) → 对应的 patch 索引
    """
    def __init__(self, img_size=224, patch_size=16):
        self.grid_size = img_size // patch_size  # 14
        self.patch_size = patch_size
        
    def select(self, gaze_points, video_tokens):
        """
        Args:
            gaze_points: (B, T, 2) 归一化凝视坐标 [0, 1]
            video_tokens: (B, T*196, D) 视频空间 token
        Returns:
            selected_tokens: (B, T, D) 凝视对应的 patch token
        """
        B, T, _ = gaze_points.shape
        # 将归一化坐标映射到 patch 索引
        grid_x = (gaze_points[..., 0] * self.grid_size).long()
        grid_y = (gaze_points[..., 1] * self.grid_size).long()
        grid_x = grid_x.clamp(0, self.grid_size - 1)
        grid_y = grid_y.clamp(0, self.grid_size - 1)
        
        # patch 索引 = y * grid_size + x
        patch_idx = grid_y * self.grid_size + grid_x  # (B, T)
        
        selected = []
        for b in range(B):
            # video_tokens[b] shape: (T*196, D)
            # 第 t 帧的 patch 索引: t*196 + patch_idx[b, t]
            indices = patch_idx[b] + torch.arange(T, device=patch_idx.device) * 196
            selected.append(video_tokens[b, indices])
        
        return torch.stack(selected, dim=0)  # (B, T, D)


class CrossAttentionBlock(nn.Module):
    """
    凝视 token 作为 Query, 选中 video patch 作为 Key/Value
    """
    def __init__(self, dim=768, num_heads=12):
        super().__init__()
        self.cross_attn = nn.MultiheadAttention(
            embed_dim=dim,
            num_heads=num_heads,
            batch_first=True
        )
        self.norm_q = nn.LayerNorm(dim)
        self.norm_kv = nn.LayerNorm(dim)
        self.ffn = nn.Sequential(
            nn.Linear(dim, dim * 4),
            nn.GELU(),
            nn.Linear(dim * 4, dim),
            nn.Dropout(0.1)
        )
        self.norm_ffn = nn.LayerNorm(dim)
        
    def forward(self, query, key_value):
        """
        Args:
            query: (B, T, D) 凝视 tokens
            key_value: (B, T, D) 选中的 video tokens
        Returns:
            fused: (B, T, D) 融合后的表示
        """
        # 跨注意力
        attn_out, _ = self.cross_attn(
            self.norm_q(query),
            self.norm_kv(key_value),
            self.norm_kv(key_value)
        )
        x = query + attn_out  # 残差连接
        
        # FFN
        ffn_out = self.ffn(self.norm_ffn(x))
        x = x + ffn_out
        return x


class SemanticFusion(nn.Module):
    """
    堆叠多层跨注意力块
    """
    def __init__(self, dim=768, num_heads=12, num_layers=2):
        super().__init__()
        self.layers = nn.ModuleList([
            CrossAttentionBlock(dim, num_heads) 
            for _ in range(num_layers)
        ])
        
    def forward(self, gaze_tokens, selected_video_tokens):
        x = gaze_tokens
        for layer in self.layers:
            x = layer(x, selected_video_tokens)
        return x  # (B, T, D)

"""
EyeCue 完整训练复现

论文：EyeCue: Driver Cognitive Distraction Detection via 
      Gaze-Empowered Egocentric Video Understanding
会议：IJCAI 2026
"""

import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from models.video_encoder import VideoEncoder
from models.gaze_encoder import GazeEncoder
from models.semantic import SemanticFusion, GazeGuidedPatchSelection
from models.head import ClassificationHead
from data.dataset import VideoGazeDataset
import argparse
import os
from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score, confusion_matrix

class EyeCue(nn.Module):
    """
    EyeCue 完整模型
    
    组合: Video Encoder + Gaze Encoder + GDSQ + Semantic Fusion + Classification Head
    """
    def __init__(self, clip_len=8, num_classes=2):
        super().__init__()
        self.clip_len = clip_len
        
        # 视频编码器: TimeSformer
        self.video_encoder = VideoEncoder(clip_len=clip_len)
        
        # 凝视编码器: 轻量 Transformer
        self.gaze_encoder = GazeEncoder(embed_dim=768, num_heads=12, num_layers=4)
        
        # 凝视引导 Patch 选择
        self.patch_selector = GazeGuidedPatchSelection(
            img_size=224, patch_size=16
        )
        
        # 语义跨注意力融合
        self.semantic = SemanticFusion(dim=768, num_heads=12, num_layers=2)
        
        # 分类头
        self.head = ClassificationHead(dim=768, num_classes=num_classes)
        
    def forward(self, video_frames, gaze_sequence):
        """
        Args:
            video_frames: (B, C, T, H, W) 视频帧序列
            gaze_sequence: (B, T, 2) 归一化凝视坐标 [0, 1]
        Returns:
            logits: (B, 2) 分类输出
        """
        # 1. 视频编码
        cls_token, video_tokens = self.video_encoder(video_frames)
        # video_tokens: (B, T*196, D)
        
        # 2. 凝视编码
        gaze_cls, gaze_tokens = self.gaze_encoder(gaze_sequence)
        # gaze_tokens: (B, T, D)
        
        # 3. 凝视引导 Patch 选择
        selected_tokens = self.patch_selector.select(
            gaze_sequence, video_tokens
        )
        # selected_tokens: (B, T, D)
        
        # 4. 语义跨注意力融合
        fused = self.semantic(gaze_tokens, selected_tokens)
        # fused: (B, T, D)
        
        # 5. 分类
        # 使用 CLS token 和融合特征
        combined = torch.cat([gaze_cls.squeeze(1), fused.mean(dim=1)], dim=-1)
        logits = self.head(combined)
        
        return logits


def train_one_epoch(model, loader, optimizer, criterion, device):
    model.train()
    total_loss = 0
    all_preds, all_labels = [], []
    
    for video, gaze, labels in loader:
        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()
        
        total_loss += loss.item()
        preds = logits.argmax(dim=-1).cpu().numpy()
        all_preds.extend(preds)
        all_labels.extend(labels.cpu().numpy())
    
    acc = accuracy_score(all_labels, all_preds)
    return total_loss / len(loader), acc


def evaluate(model, loader, device):
    model.eval()
    all_preds, all_labels, all_probs = [], [], []
    
    with torch.no_grad():
        for video, gaze, labels in loader:
            video = video.to(device)
            gaze = gaze.to(device)
            logits = model(video, gaze)
            probs = torch.softmax(logits, dim=-1)
            preds = logits.argmax(dim=-1)
            
            all_preds.extend(preds.cpu().numpy())
            all_labels.extend(labels.numpy())
            all_probs.extend(probs[:, 1].cpu().numpy())  # 分心概率
    
    acc = accuracy_score(all_labels, all_preds)
    prec = precision_score(all_labels, all_preds)
    rec = recall_score(all_labels, all_preds)
    f1 = f1_score(all_labels, all_preds)
    cm = confusion_matrix(all_labels, all_preds)
    
    return {
        'accuracy': acc,
        'precision': prec,
        'recall': rec,
        'f1': f1,
        'confusion_matrix': cm,
        'probabilities': all_probs
    }


if __name__ == "__main__":
    # 参数配置（按论文设置）
    config = {
        'clip_len': 8,        # 每个片段采样 8 帧
        'batch_size': 4,      # 论文设置
        'epochs': 15,
        'lr': 1e-5,
        'device': 'cuda' if torch.cuda.is_available() else 'cpu',
        'train_list': '/path/to/train.txt',
        'val_list': '/path/to/val.txt',
        'save_dir': 'checkpoints'
    }
    
    os.makedirs(config['save_dir'], exist_ok=True)
    
    # 数据集
    train_ds = VideoGazeDataset(
        config['train_list'], clip_len=config['clip_len']
    )
    val_ds = VideoGazeDataset(
        config['val_list'], clip_len=config['clip_len']
    )
    
    train_loader = DataLoader(
        train_ds, batch_size=config['batch_size'], 
        shuffle=True, num_workers=1
    )
    val_loader = DataLoader(
        val_ds, batch_size=config['batch_size'],
        num_workers=1
    )
    
    # 模型
    model = EyeCue(clip_len=config['clip_len']).to(config['device'])
    optimizer = torch.optim.AdamW(model.parameters(), lr=config['lr'])
    criterion = nn.CrossEntropyLoss()
    
    # 训练
    best_acc = 0
    for epoch in range(config['epochs']):
        train_loss, train_acc = train_one_epoch(
            model, train_loader, optimizer, criterion, config['device']
        )
        val_metrics = evaluate(model, val_loader, config['device'])
        
        print(f"Epoch {epoch+1}/{config['epochs']}")
        print(f"  Train Loss: {train_loss:.4f}, Train Acc: {train_acc:.4f}")
        print(f"  Val Acc: {val_metrics['accuracy']:.4f}, "
              f"Val F1: {val_metrics['f1']:.4f}")
        
        if val_metrics['accuracy'] > best_acc:
            best_acc = val_metrics['accuracy']
            torch.save(model.state_dict(), 
                      f"{config['save_dir']}/best_model.pth")
            print(f"  ✅ 新最佳模型已保存 (Acc: {best_acc:.4f})")
    
    print(f"\n训练完成。最佳准确率: {best_acc:.4f}")
    
    # 最终评估
    model.load_state_dict(
        torch.load(f"{config['save_dir']}/best_model.pth")
    )
    final_metrics = evaluate(model, val_loader, config['device'])
    
    print("\n=== 最终评估结果 ===")
    print(f"准确率: {final_metrics['accuracy']:.4f}")
    print(f"精确率: {final_metrics['precision']:.4f}")
    print(f"召回率: {final_metrics['recall']:.4f}")
    print(f"F1: {final_metrics['f1']:.4f}")
    print(f"混淆矩阵:\n{final_metrics['confusion_matrix']}")

"""
EyeCue 批量推理与评估

用法: python infer.py --test_list test.txt --checkpoint checkpoints/best_model.pth
"""

import csv
import torch
from torch.utils.data import DataLoader
from models.video_encoder import VideoEncoder
from models.gaze_encoder import GazeEncoder
from models.semantic import SemanticFusion, GazeGuidedPatchSelection
from models.head import ClassificationHead
from data.dataset import VideoGazeDataset
from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score, confusion_matrix
import argparse

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--test_list', required=True)
    parser.add_argument('--checkpoint', default='checkpoints/best_model.pth')
    parser.add_argument('--out_csv', default='predictions.csv')
    parser.add_argument('--batch_size', type=int, default=4)
    parser.add_argument('--clip_len', type=int, default=8)
    parser.add_argument('--seed', type=int, default=42)
    parser.add_argument('--device', default='auto')
    args = parser.parse_args()
    
    device = ('cuda' if args.device == 'auto' and torch.cuda.is_available() 
              else args.device)
    if device == 'auto':
        device = 'cpu'
    
    torch.manual_seed(args.seed)
    
    # 加载模型
    model = EyeCue(clip_len=args.clip_len).to(device)
    model.load_state_dict(torch.load(args.checkpoint, map_location=device))
    model.eval()
    
    # 数据
    test_ds = VideoGazeDataset(args.test_list, clip_len=args.clip_len)
    test_loader = DataLoader(test_ds, batch_size=args.batch_size, num_workers=1)
    
    all_preds, all_labels, filenames = [], [], []
    
    with torch.no_grad():
        for video, gaze, labels in test_loader:
            video = video.to(device)
            gaze = gaze.to(device)
            logits = model(video, gaze)
            probs = torch.softmax(logits, dim=-1)
            preds = logits.argmax(dim=-1)
            
            all_preds.extend(preds.cpu().numpy())
            all_labels.extend(labels.numpy())
    
    # 指标
    acc = accuracy_score(all_labels, all_preds)
    prec = precision_score(all_labels, all_preds)
    rec = recall_score(all_labels, all_preds)
    f1 = f1_score(all_labels, all_preds)
    cm = confusion_matrix(all_labels, all_preds)
    
    print(f"Accuracy: {acc:.4f}")
    print(f"Precision: {prec:.4f}")
    print(f"Recall: {rec:.4f}")
    print(f"F1: {f1:.4f}")
    print(f"Confusion Matrix:\n{cm}")
    
    # 保存预测
    with open(args.out_csv, 'w', newline='') as f:
        writer = csv.writer(f)
        writer.writerow(['filename', 'label', 'pred', 'prob_distracted'])
        for i, (fname, label, pred) in enumerate(zip(filenames, all_labels, all_preds)):
            prob = all_probs[i] if 'all_probs' in dir() else 0.0
            writer.writerow([fname, label, pred, f"{prob:.4f}"])

if __name__ == "__main__":
    main()