"""
GazeFS 视线轨迹预测与稳定化

核心发现：
1. 目标获取产生任务对齐的视线-头部动力学变化
2. 视线轨迹保留持久的目标相对残余方向
3. 历史窗口改善相位解码优于当前端点

两个阶段：
- Search（搜索）：视线扫描寻找目标
- Focus（聚焦）：视线锁定目标并稳定

IMS 应用：
- 搜索状态 → 分心检测（视线频繁扫描 = 注意力不集中）
- 聚焦状态 → 正常驾驶（视线锁定道路 = 集中注意力）
- 搜索时长 → 分心程度的量化指标
"""

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

class GazeHistoryEncoder(nn.Module):
    """
    视线-头部历史编码器
    
    处理变长历史序列，编码时序模式
    """
    
    def __init__(self, input_dim: int = 4, hidden_dim: int = 128):
        super().__init__()
        # 输入: [gaze_x, gaze_y, head_x, head_y]
        self.proj = nn.Linear(input_dim, hidden_dim)
        
        # LSTM 处理变长历史
        self.lstm = nn.LSTM(
            input_size=hidden_dim,
            hidden_size=hidden_dim,
            num_layers=2,
            batch_first=True,
            dropout=0.1
        )
        
    def forward(self, history: torch.Tensor, 
                lengths: torch.Tensor = None) -> torch.Tensor:
        """
        Args:
            history: (B, T, 4) 视线-头部历史序列
            lengths: (B,) 每个样本的实际长度
        Returns:
            encoded: (B, hidden_dim) 历史编码
        """
        x = self.proj(history)
        out, (h_n, c_n) = self.lstm(x)
        # 取最后一个有效输出
        if lengths is not None:
            idx = (lengths - 1).view(-1, 1, 1).expand(-1, 1, out.size(-1))
            encoded = out.gather(1, idx).squeeze(1)
        else:
            encoded = out[:, -1]
        return encoded


class GazeFSNet(nn.Module):
    """
    GazeFS: 视线轨迹预测与稳定化网络
    
    输入: 变长视线-头部历史
    输出: 
    1. 下一个目标中心方向 (gaze_x, gaze_y)
    2. 搜索/聚焦状态估计 (search_prob, focus_prob)
    
    性能 (论文报告):
    - Focus 偏差降低: 0.182°
    - P90 误差降低: 0.400°
    - 相位分类准确率: 0.925
    - AUPRC: 0.993
    """
    
    def __init__(self, input_dim: int = 4, hidden_dim: int = 128):
        super().__init__()
        self.encoder = GazeHistoryEncoder(input_dim, hidden_dim)
        
        # 方向预测头
        self.direction_head = nn.Sequential(
            nn.Linear(hidden_dim, 64),
            nn.ReLU(),
            nn.Linear(64, 2)  # (gaze_x, gaze_y)
        )
        
        # 相位分类头（搜索 vs 聚焦）
        self.phase_head = nn.Sequential(
            nn.Linear(hidden_dim, 64),
            nn.ReLU(),
            nn.Linear(64, 2),  # (search_prob, focus_prob)
            nn.Softmax(dim=-1)
        )
        
    def forward(self, history: torch.Tensor,
                lengths: torch.Tensor = None) -> dict:
        encoded = self.encoder(history, lengths)
        
        direction = self.direction_head(encoded)
        phase = self.phase_head(encoded)
        
        return {
            'direction': direction,
            'phase': phase,  # (search_prob, focus_prob)
            'encoded': encoded
        }


class GazeStabilizer:
    """
    视线稳定化后处理
    
    论文指出：GazeFS 改善目标中心和残余收缩，
    但时间平滑是独立目标
    
    IMS 实现方案：
    1. GazeFS 预测目标方向
    2. Kalman 滤波器做时间平滑
    3. 搜索/聚焦状态用于分心检测
    """
    
    def __init__(self, alpha: float = 0.3):
        """
        Args:
            alpha: 指数移动平均系数 (0-1)
                   越大越平滑，越小越跟随
        """
        self.alpha = alpha
        self.smoothed = None
        self.model = GazeFSNet()
        self.model.eval()
        
    def process(self, raw_gaze: torch.Tensor, 
                head_data: torch.Tensor,
                history: list = None) -> dict:
        """
        在线处理单帧
        
        Args:
            raw_gaze: (2,) 当前帧原始视线 [x, y]
            head_data: (2,) 当前帧头部 [yaw, pitch]
            history: 历史窗口 [(gaze, head), ...]
        
        Returns:
            稳定化后的视线 + 相位状态
        """
        # 构建历史输入
        if history is None or len(history) == 0:
            # 无历史时直接返回原始值
            return {
                'gaze': raw_gaze,
                'phase': 'unknown',
                'search_prob': 0.5
            }
        
        # 准备模型输入 (B=1, T, 4)
        hist_tensor = torch.tensor(
            [[g[0] + h[0], g[1] + h[1]] for g, h in history[-30:]],
            dtype=torch.float32
        ).unsqueeze(0)
        
        with torch.no_grad():
            out = self.model(hist_tensor)
        
        predicted_dir = out['direction'][0]
        phase = out['phase'][0]
        search_prob = phase[0].item()
        
        # EMA 平滑
        if self.smoothed is None:
            self.smoothed = predicted_dir.clone()
        else:
            self.smoothed = self.alpha * self.smoothed + (1 - self.alpha) * predicted_dir
        
        return {
            'gaze': self.smoothed,
            'phase': 'search' if search_prob > 0.5 else 'focus',
            'search_prob': search_prob,
            'raw_gaze': raw_gaze,
        }


# ==================== 测试代码 ====================
if __name__ == "__main__":
    import time
    import numpy as np
    
    stabilizer = GazeStabilizer(alpha=0.3)
    
    # 模拟 5 秒视线数据 (30fps)
    frames = 150
    # 前 2 秒：搜索阶段（视线频繁跳转）
    # 后 3 秒：聚焦阶段（视线稳定）
    
    search_gaze = np.array([
        np.random.uniform(-0.5, 0.5, 2) for _ in range(60)
    ])
    focus_gaze = np.array([
        np.random.normal(0, 0.02, 2) for _ in range(90)
    ])
    
    head_data = np.array([
        np.random.normal(0, 0.01, 2) for _ in range(frames)
    ])
    
    all_gaze = np.vstack([search_gaze, focus_gaze])
    
    # 在线处理
    history = []
    search_count = 0
    focus_count = 0
    
    print("=" * 60)
    print("GazeFS 视线轨迹预测与稳定化测试")
    print("=" * 60)
    
    start = time.time()
    for i in range(frames):
        raw = torch.tensor(all_gaze[i], dtype=torch.float32)
        head = torch.tensor(head_data[i], dtype=torch.float32)
        
        result = stabilizer.process(raw, head, history)
        history.append((raw.numpy(), head.numpy()))
        
        if result['phase'] == 'search':
            search_count += 1
        else:
            focus_count += 1
        
        # 打印关键时间点
        if i == 59:
            print(f"\n[Frame 60] 搜索→聚焦转换点")
            print(f"  搜索帧: {search_count}, 聚焦帧: {focus_count}")
            print(f"  当前相位: {result['phase']} (search_prob={result['search_prob']:.3f})")
    
    total_time = (time.time() - start) * 1000
    print(f"\n总处理时间: {total_time:.1f} ms ({total_time/frames:.2f} ms/frame)")
    print(f"搜索帧: {search_count}/{frames} ({search_count/frames*100:.0f}%)")
    print(f"聚焦帧: {focus_count}/{frames} ({focus_count/frames*100:.0f}%)")
    print(f"相位分类准确率: {0.925 * 100:.1f}% (论文报告)")
    print(f"P90 误差降低: 0.400° (论文报告)")