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

class DifferentialGazeModel(nn.Module):
    """
    差分视线估计模型
    
    论文: arXiv:2610.11036
    核心: 预测两帧之间的视线位移而非绝对视线
    
    架构:
    1. 权重共享ResNet-18编码两帧眼部图像
    2. MLP映射编码差到2D视线位移
    3. 绝对视线 = 锚点视线 + 预测位移
    
    优势:
    - 消除个体化偏移（差分操作平移不变）
    - 对标定漂移鲁棒（相对量不受绝对漂移影响）
    - 在窄范围中区分微小变化更敏感
    """
    
    def __init__(self, hidden_dim: int = 256):
        super().__init__()
        
        # 权重共享的眼部图像编码器
        import torchvision.models as models
        resnet = models.resnet18(weights='DEFAULT')
        resnet.fc = nn.Identity()
        self.encoder = resnet  # 输出512维
        
        # 差分映射头
        self.diff_head = nn.Sequential(
            nn.Linear(512 * 2, hidden_dim),
            nn.BatchNorm1d(hidden_dim),
            nn.ReLU(inplace=True),
            nn.Dropout(0.2),
            nn.Linear(hidden_dim, hidden_dim // 2),
            nn.BatchNorm1d(hidden_dim // 2),
            nn.ReLU(inplace=True),
            nn.Dropout(0.2),
            nn.Linear(hidden_dim // 2, 2)  # (Δpitch, Δyaw) in radians
        )
        
        # 时序上下文模块（可选）
        self.temporal_lstm = nn.LSTM(
            input_size=2,  # 差分序列
            hidden_size=64,
            num_layers=1,
            batch_first=True
        )
        self.temporal_head = nn.Linear(64, 2)
        
    def forward(
        self,
        anchor_eye: torch.Tensor,   # (B, 3, 64, 64) 锚点帧眼部图像
        current_eye: torch.Tensor,   # (B, 3, 64, 64) 当前帧眼部图像
        anchor_gaze: torch.Tensor,   # (B, 2) 锚点视线 (pitch, yaw)
        prev_diffs: torch.Tensor = None  # (B, T, 2) 过去T帧差分序列
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        前向传播
        
        Args:
            anchor_eye: 标定锚点帧的眼部裁剪
            current_eye: 当前帧的眼部裁剪
            anchor_gaze: 锚点的已知视线
            prev_diffs: 过去T帧的差分序列（用于时序增强）
        
        Returns:
            predicted_gaze: (B, 2) 预测的绝对视线
            predicted_diff: (B, 2) 预测的差分
        """
        # 编码两帧
        e_anchor = self.encoder(anchor_eye)   # (B, 512)
        e_current = self.encoder(current_eye) # (B, 512)
        
        # 拼接并预测差分
        concat = torch.cat([e_anchor, e_current], dim=-1)  # (B, 1024)
        diff = self.diff_head(concat)  # (B, 2) 视线位移
        
        # 绝对视线 = 锚点 + 差分
        predicted_gaze = anchor_gaze + diff
        
        # 时序增强（如果有历史差分）
        if prev_diffs is not None:
            temporal_out, _ = self.temporal_lstm(prev_diffs)
            temporal_correction = self.temporal_head(temporal_out[:, -1, :])
            predicted_gaze = predicted_gaze + temporal_correction
        
        return predicted_gaze, diff
    
    def calibrate(self, anchor_eye, anchor_gaze):
        """
        被试标定：记录锚点帧和视线
        
        在长时间场景中，需要周期性重新标定
        因为标定漂移会累积
        """
        self.register_buffer('anchor_eye', anchor_eye)
        self.register_buffer('anchor_gaze', anchor_gaze)


class LongDurationGazeTracker:
    """
    长时驾驶视线追踪系统
    
    集成:
    1. 差分视线估计
    2. 周期性重标定（对抗漂移）
    3. 时序平滑
    4. 注视/扫视事件检测
    """
    
    def __init__(self, model, recalibration_interval=300):
        """
        Args:
            model: DifferentialGazeModel
            recalibration_interval: 重标定间隔（帧数，300帧=10秒@30fps）
        """
        self.model = model
        self.recal_interval = recalibration_interval
        self.frame_count = 0
        self.gaze_history = []
        self.diff_history = []
        self.anchor = None
        
    def initialize(self, anchor_eye, anchor_gaze):
        """初始化锚点"""
        self.model.calibrate(anchor_eye, anchor_gaze)
        self.anchor = (anchor_eye, anchor_gaze)
        self.frame_count = 0
        self.gaze_history = [anchor_gaze]
        self.diff_history = []
    
    def step(self, current_eye):
        """
        单帧处理
        
        Returns:
            gaze: (2,) 预测视线
            is_fixation: bool 是否为注视
            is_saccade: bool 是否为扫视
            need_recalibration: bool 是否需要重标定
        """
        if self.anchor is None:
            return None, False, False, True
        
        anchor_eye, anchor_gaze = self.anchor
        
        # 准备历史差分
        prev_diffs = None
        if len(self.diff_history) > 5:
            import numpy as np
            prev_diffs = torch.stack([
                torch.tensor(d) for d in self.diff_history[-10:]
            ]).unsqueeze(0)
        
        # 预测
        with torch.no_grad():
            gaze, diff = self.model(
                anchor_eye.unsqueeze(0),
                current_eye.unsqueeze(0),
                anchor_gaze.unsqueeze(0),
                prev_diffs
            )
        
        gaze = gaze[0]
        diff = diff[0]
        
        # 更新历史
        self.gaze_history.append(gaze)
        self.diff_history.append(diff)
        self.frame_count += 1
        
        # 注视/扫射检测
        diff_mag = diff.norm().item()
        is_saccade = diff_mag > 0.05  # >~3° 突变
        is_fixation = diff_mag < 0.01  # <~0.6° 微动
        
        # 重标定检查
        need_recal = self.frame_count >= self.recal_interval
        
        return gaze.numpy(), is_fixation, is_saccade, need_recal


# 测试
if __name__ == "__main__":
    print("=" * 60)
    print("差分视线估计 - 长时窄范围场景测试")
    print("=" * 60)
    
    model = DifferentialGazeModel()
    total_params = sum(p.numel() for p in model.parameters())
    print(f"\n模型参数量: {total_params:,}")
    
    # 模拟锚点和当前帧
    B = 2
    anchor_eye = torch.randn(B, 3, 64, 64)
    current_eye = torch.randn(B, 3, 64, 64)
    anchor_gaze = torch.tensor([[0.0, 0.0], [0.1, -0.05]])  # radians
    
    # 预测
    gaze, diff = model(anchor_eye, current_eye, anchor_gaze)
    
    print(f"\n输入:")
    print(f"  锚点视线: {anchor_gaze}")
    print(f"  锚点眼部: {anchor_eye.shape}")
    print(f"  当前眼部: {current_eye.shape}")
    print(f"\n输出:")
    print(f"  预测视线: {gaze}")
    print(f"  视线差分: {diff}")
    print(f"  差分幅度(rad): {diff.norm(dim=-1)}")
    print(f"  差分幅度(°): {diff.norm(dim=-1) * 180 / 3.14159}")
    
    # 长时追踪系统测试
    print(f"\n{'='*60}")
    print("长时追踪系统测试（10帧）")
    print(f"{'='*60}")
    
    tracker = LongDurationGazeTracker(model, recalibration_interval=5)
    
    # 初始化
    anchor_eye_single = torch.randn(3, 64, 64)
    anchor_gaze_single = torch.tensor([0.0, 0.0])
    tracker.initialize(anchor_eye_single, anchor_gaze_single)
    
    print(f"\n{'帧':>4} | {'pitch(°)':>8} | {'yaw(°)':>8} | {'Δ(°)':>6} | {'状态':>8} | {'重标定':>6}")
    print("-" * 55)
    
    import numpy as np
    for i in range(10):
        current = torch.randn(3, 64, 64)
        gaze, is_fix, is_sac, need_recal = tracker.step(current)
        
        if gaze is not None:
            pitch = gaze[0] * 180 / np.pi
            yaw = gaze[1] * 180 / np.pi
            diff_mag = np.sqrt(gaze[0]**2 + gaze[1]**2) * 180 / np.pi
            status = "注视" if is_fix else ("扫视" if is_sac else "追踪")
            recal = "是" if need_recal else "否"
            print(f"{i:4d} | {pitch:8.2f} | {yaw:8.2f} | {diff_mag:6.2f} | {status:>8} | {recal:>6}")
    
    print(f"\n历史长度: {len(tracker.gaze_history)}")

# 重标定策略建议
recalibration_strategies = {
    "时间触发": {
        "间隔": "每10分钟",
        "触发": "自动",
        "代价": "1帧标定（驾驶员不知情）"
    },
    "事件触发": {
        "条件": "追踪器置信度高+面部正面",
        "优势": "不影响驾驶体验",
        "实现": "质量评估器自动判断"
    },
    "混合策略": {
        "默认": "时间触发10分钟",
        "覆盖": "事件触发加速",
        "推荐": True
    }
}