import torch
import torch.nn as nn

class CausalGazeForecaster(nn.Module):
    """
    Context-Aware Causal Gaze Forecaster
    
    在追踪丢失期间:
    1. 分析丢失前的视线轨迹
    2. 结合驾驶上下文（车辆状态、场景）
    3. 因果推断丢失期间的视线方向
    
    关键假设: 驾驶员的视线转移由驾驶意图因果驱动
    """
    
    def __init__(self, feat_dim: int = 128, n_context: int = 30):
        super().__init__()
        
        # 视线历史编码器
        self.gaze_encoder = nn.LSTM(
            input_size=2,  # (pitch, yaw)
            hidden_size=feat_dim,
            num_layers=2,
            batch_first=True,
            bidirectional=True
        )
        
        # 驾驶上下文编码器
        self.context_encoder = nn.Sequential(
            nn.Linear(10, 32), nn.ReLU(),  # 车辆状态: 速度/转向/车道/制动
            nn.Linear(32, feat_dim)
        )
        
        # 因果推断网络
        self.causal_net = nn.TransformerDecoder(
            nn.TransformerDecoderLayer(
                d_model=feat_dim * 2,
                nhead=8,
                dim_feedforward=256,
                batch_first=True
            ),
            num_layers=3
        )
        
        # 视线预测头
        self.forecast_head = nn.Sequential(
            nn.Linear(feat_dim * 2, 64), nn.ReLU(),
            nn.Linear(64, 2)  # (pitch, yaw)
        )
        
        # 置信度预测
        self.confidence_head = nn.Sequential(
            nn.Linear(feat_dim * 2, 32), nn.ReLU(),
            nn.Linear(32, 1), nn.Sigmoid()
        )
    
    def forward(
        self,
        gaze_history: torch.Tensor,  # (B, T_hist, 2) 丢失前视线
        context: torch.Tensor,         # (B, 10) 驾驶上下文
        n_forecast: int = 30           # 预测多少帧
    ) -> dict:
        """
        预测追踪丢失期间的视线方向
        
        Returns:
            forecast_gaze: (B, n_forecast, 2) 预测视线
            confidence: (B, n_forecast) 每帧置信度
        """
        # 编码视线历史
        _, (h_n, _) = self.gaze_encoder(gaze_history)
        # 拼接双向LSTM最后隐状态
        gaze_feat = torch.cat([h_n[-2], h_n[-1]], dim=1)  # (B, 2*feat_dim)
        
        # 编码驾驶上下文
        ctx_feat = self.context_encoder(context)  # (B, feat_dim)
        
        # 融合
        combined = gaze_feat + torch.cat([ctx_feat, ctx_feat], dim=1)
        
        # 生成查询序列
        queries = combined.unsqueeze(1).expand(-1, n_forecast, -1)
        
        # 因果解码
        decoded = self.causal_net(
            queries,
            combined.unsqueeze(1)  # 记忆
        )
        
        # 预测
        forecast_gaze = self.forecast_head(decoded)
        confidence = self.confidence_head(decoded).squeeze(-1)
        
        return {
            'forecast_gaze': forecast_gaze,
            'confidence': confidence
        }


# 测试
if __name__ == "__main__":
    model = CausalGazeForecaster(feat_dim=128, n_context=30)
    
    # 模拟：30帧视线历史 + 驾驶上下文
    gaze_hist = torch.randn(4, 30, 2)  # 1s@30fps
    ctx = torch.randn(4, 10)           # 速度+转向+车道+...
    
    out = model(gaze_hist, ctx, n_forecast=45)  # 预测1.5s
    
    print(f"视线历史: {gaze_hist.shape}")
    print(f"预测视线: {out['forecast_gaze'].shape}")
    print(f"置信度: {out['confidence'].shape}")
    print(f"平均置信度: {out['confidence'].mean():.3f}")