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

class CausalContextGatedForecaster(nn.Module):
    """
    CCGF: 因果上下文门控预测器
    
    论文: arXiv:2609.12374
    核心功能: 在追踪器丢失期间，仅用过去数据预测驾驶员视线
    
    架构:
    1. 历史LSTM: 编码60帧预丢失视线+头姿
    2. 场景LSTM: 编码DINOv3场景特征
    3. 可靠性门控: 动态融合历史和场景
    4. 空间概率解码: 24×32分布 → 质心
    """
    
    def __init__(
        self,
        gaze_dim: int = 2,        # 视线 (x, y) 像素坐标
        head_dim: int = 6,        # 头姿 (pitch, yaw, roll + velocity)
        scene_dim: int = 768,     # DINOv3特征维度
        hidden_dim: int = 256,
        grid_h: int = 24,
        grid_w: int = 32,
        history_len: int = 60
    ):
        super().__init__()
        self.grid_h = grid_h
        self.grid_w = grid_w
        self.history_len = history_len
        
        # 历史编码器: 2层LSTM
        self.history_lstm = nn.LSTM(
            input_size=gaze_dim + head_dim,
            hidden_size=hidden_dim,
            num_layers=2,
            batch_first=True,
            dropout=0.1
        )
        
        # 场景编码器: 1层LSTM
        self.scene_lstm = nn.LSTM(
            input_size=scene_dim,
            hidden_size=hidden_dim,
            num_layers=1,
            batch_first=True
        )
        
        # 可靠性门控
        # 输入: 当前时间步、历史隐状态、场景隐状态
        self.gate = nn.Sequential(
            nn.Linear(hidden_dim * 2 + 1, hidden_dim),
            nn.GELU(),
            nn.Linear(hidden_dim, 1),
            nn.Sigmoid()
        )
        
        # 空间概率解码器
        self.decoder = nn.Sequential(
            nn.Linear(hidden_dim, hidden_dim * 2),
            nn.GELU(),
            nn.Linear(hidden_dim * 2, grid_h * grid_w)
        )
        
        # 位置编码
        self.register_buffer(
            'grid_coords',
            self._create_grid_coords(grid_h, grid_w)
        )
    
    def _create_grid_coords(self, h: int, w: int) -> torch.Tensor:
        """创建空间网格坐标"""
        ys = torch.linspace(0, 1, h)
        xs = torch.linspace(0, 1, w)
        grid = torch.stack(torch.meshgrid(ys, xs, indexing='ij'))
        return grid.reshape(2, -1).T  # (h*w, 2)
    
    def forward(
        self,
        gaze_history: torch.Tensor,    # (B, T, 2)
        head_history: torch.Tensor,     # (B, T, 6)
        scene_features: torch.Tensor,  # (B, T, 768)
        dropout_progress: float = 0.0  # 丢失进度 0-1
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        前向传播
        
        Args:
            gaze_history: 过去T帧视线坐标 (B, T, 2)
            head_history: 过去T帧头姿 (B, T, 6)
            scene_features: 场景DINOv3特征 (B, T, 768)
            dropout_progress: 当前在丢失期间的进度 [0, 1]
        
        Returns:
            predicted_gaze: (B, 2) 预测视线位置
            prob_map: (B, 24, 32) 空间概率分布
        """
        B = gaze_history.shape[0]
        
        # 历史编码
        hist_input = torch.cat([gaze_history, head_history], dim=-1)
        _, (h_n, _) = self.history_lstm(hist_input)
        hist_hidden = h_n[-1]  # (B, hidden_dim) 最后一层
        
        # 场景编码
        _, (s_n, _) = self.scene_lstm(scene_features)
        scene_hidden = s_n[-1]  # (B, hidden_dim)
        
        # 可靠性门控
        progress = torch.tensor(dropout_progress, 
                                device=gaze_history.device).expand(B, 1)
        gate_input = torch.cat([hist_hidden, scene_hidden, progress], dim=-1)
        alpha = self.gate(gate_input)  # (B, 1)
        
        # 融合
        fused = alpha * hist_hidden + (1 - alpha) * scene_hidden
        
        # 解码为空间概率分布
        logits = self.decoder(fused)  # (B, h*w)
        prob_map = torch.softmax(logits, dim=-1)  # (B, h*w)
        prob_map = prob_map.reshape(B, self.grid_h, self.grid_w)
        
        # 质心 = 预测视线位置
        grid = self.grid_coords.to(gaze_history.device)  # (h*w, 2)
        prob_flat = prob_map.reshape(B, -1)  # (B, h*w)
        predicted_gaze = (prob_flat.unsqueeze(-1) * grid).sum(dim=1)  # (B, 2)
        
        return predicted_gaze, prob_map


class DropoutRecoverySystem:
    """
    CCGF在线丢失恢复系统
    
    集成到DMS中的完整pipeline:
    1. 检测追踪器丢失
    2. 启动CCGF预测
    3. 追踪器恢复后平滑切换
    """
    
    def __init__(self, history_len=60, max_dropout_len=120):
        self.forecaster = CausalContextGatedForecaster(history_len=history_len)
        self.history_len = history_len
        self.max_dropout_len = max_dropout_len
        
        # 历史缓冲区
        self.gaze_buffer = []
        self.head_buffer = []
        self.scene_buffer = []
        
        # 状态
        self.is_dropout = False
        self.dropout_frame = 0
        
    def update(self, gaze, head_pose, scene_feat, tracker_status):
        """
        每帧调用
        
        Args:
            gaze: (2,) 视线坐标, None if tracker lost
            head_pose: (6,) 头姿
            scene_feat: (768,) 场景特征
            tracker_status: 'tracking' or 'lost'
        
        Returns:
            predicted_gaze: (2,) 视线坐标（真实或预测）
            is_predicted: bool 是否为预测值
        """
        if tracker_status == 'tracking':
            # 正常追踪：更新缓冲区
            self.gaze_buffer.append(gaze.copy())
            self.head_buffer.append(head_pose.copy())
            self.scene_buffer.append(scene_feat.copy())
            
            # 保持窗口大小
            if len(self.gaze_buffer) > self.history_len:
                self.gaze_buffer.pop(0)
                self.head_buffer.pop(0)
                self.scene_buffer.pop(0)
            
            self.is_dropout = False
            self.dropout_frame = 0
            return gaze, False
            
        else:
            # 追踪丢失：启动CCGF预测
            if not self.is_dropout:
                self.is_dropout = True
                self.dropout_frame = 0
            
            self.dropout_frame += 1
            
            # 需要足够的历史
            if len(self.gaze_buffer) < self.history_len:
                return None, False
            
            # 准备输入
            gaze_hist = torch.tensor(
                np.array(self.gaze_buffer)
            ).unsqueeze(0).float()
            head_hist = torch.tensor(
                np.array(self.head_buffer)
            ).unsqueeze(0).float()
            scene_hist = torch.tensor(
                np.array(self.scene_buffer)
            ).unsqueeze(0).float()
            
            # 预测
            progress = min(self.dropout_frame / self.max_dropout_len, 1.0)
            with torch.no_grad():
                pred_gaze, prob_map = self.forecaster(
                    gaze_hist, head_hist, scene_hist,
                    dropout_progress=progress
                )
            
            return pred_gaze[0].numpy(), True


import numpy as np

# 测试
if __name__ == "__main__":
    print("=" * 60)
    print("CCGF 因果视线恢复 - 测试")
    print("=" * 60)
    
    # 模型初始化
    model = CausalContextGatedForecaster()
    
    total_params = sum(p.numel() for p in model.parameters())
    print(f"\n模型参数量: {total_params:,}")
    
    # 模拟输入
    B, T = 1, 60
    gaze_hist = torch.randn(B, T, 2)
    head_hist = torch.randn(B, T, 6)
    scene_feat = torch.randn(B, T, 768)
    
    # 预测
    pred_gaze, prob_map = model(gaze_hist, head_hist, scene_feat, 0.3)
    
    print(f"\n输入:")
    print(f"  视线历史: {gaze_hist.shape}")
    print(f"  头姿历史: {head_hist.shape}")
    print(f"  场景特征: {scene_feat.shape}")
    print(f"\n输出:")
    print(f"  预测视线: {pred_gaze.shape} → {pred_gaze[0].detach().numpy()}")
    print(f"  概率分布: {prob_map.shape}")
    print(f"  概率和: {prob_map[0].sum().item():.4f}")
    
    # 完整系统测试
    print(f"\n{'='*60}")
    print("在线丢失恢复系统测试")
    print(f"{'='*60}")
    
    system = DropoutRecoverySystem(history_len=60)
    
    # 填充历史
    for i in range(60):
        gaze = np.array([0.5 + 0.01*i, 0.5])
        head = np.random.randn(6)
        scene = np.random.randn(768)
        system.update(gaze, head, scene, 'tracking')
    
    print(f"历史缓冲区已填充: {len(system.gaze_buffer)}帧")
    
    # 模拟追踪丢失
    print(f"\n模拟追踪丢失（5帧）:")
    print(f"{'帧':>4} | {'x':>8} | {'y':>8} | {'预测?':>6}")
    print("-" * 36)
    
    for i in range(5):
        gaze = system.update(
            None, np.random.randn(6), np.random.randn(768), 'lost'
        )
        if gaze[0] is not None:
            x, y = gaze[0]
            print(f"{i:4d} | {x:8.4f} | {y:8.4f} | {'是' if gaze[1] else '否':>6}")