import torch
import torch.nn as nn
import torch.nn.functional as F

class AttentiveGazeModel(nn.Module):
    """
    AttentiveGaze: 不确定性感知的多模态视线估计模型
    
    三大创新：
    1. 注意力增强眼部特征提取（空间+通道注意力）
    2. 跨模态注意力融合 + 可学习门控
    3. 不确定性量化（异方差回归）
    """
    
    def __init__(self, eye_backbone='resnet18', 
                 face_backbone='resnet18',
                 feature_dim=256,
                 num_heads=4):
        super().__init__()
        
        # 眼部特征提取器（注意力增强）
        self.eye_encoder = AttentionEyeEncoder(
            backbone=eye_backbone,
            feature_dim=feature_dim
        )
        
        # 面部上下文编码器
        self.face_encoder = FaceContextEncoder(
            backbone=face_backbone,
            feature_dim=feature_dim
        )
        
        # 头部姿态编码器
        self.head_pose_encoder = HeadPoseEncoder(
            feature_dim=64
        )
        
        # 跨模态注意力融合
        self.cross_modal_fusion = CrossModalAttentionFusion(
            eye_dim=feature_dim,
            face_dim=feature_dim,
            head_dim=64,
            hidden_dim=256,
            num_heads=num_heads
        )
        
        # 多头嵌入
        self.multi_head_embedding = MultiHeadEmbedding(
            input_dim=256,
            num_heads=num_heads,
            head_dim=64
        )
        
        # 视线回归头 + 不确定性
        self.gaze_regressor = UncertaintyGazeRegressor(
            input_dim=num_heads * 64,
            output_dim=2  # (pitch, yaw)
        )
    
    def forward(self, eye_img: torch.Tensor, 
                face_img: torch.Tensor,
                head_pose: torch.Tensor) -> dict:
        """
        Args:
            eye_img: (B, 3, 64, 64) 眼部裁剪
            face_img: (B, 3, 224, 224) 全脸图像
            head_pose: (B, 3) 头部姿态 (pitch, yaw, roll)
        
        Returns:
            {'gaze': (B, 2), 'uncertainty': (B, 1)}
        """
        # 各模态特征提取
        eye_feat = self.eye_encoder(eye_img)      # (B, feature_dim)
        face_feat = self.face_encoder(face_img)     # (B, feature_dim)
        head_feat = self.head_pose_encoder(head_pose)  # (B, 64)
        
        # 跨模态融合
        fused = self.cross_modal_fusion(
            eye_feat, face_feat, head_feat
        )  # (B, 256)
        
        # 多头嵌入
        embedded = self.multi_head_embedding(fused)  # (B, num_heads*64)
        
        # 视线 + 不确定性
        result = self.gaze_regressor(embedded)
        
        return result


class AttentionEyeEncoder(nn.Module):
    """
    注意力增强眼部特征提取器
    
    空间注意力：学习眼部图像中哪些位置最重要
    通道注意力：学习哪些特征通道最相关
    
    关键：在遮挡/光照差时，注意力可降权受损区域
    """
    
    def __init__(self, backbone='resnet18', feature_dim=256):
        super().__init__()
        
        # 骨干网络（预训练）
        from torchvision.models import resnet18
        self.backbone = resnet18(pretrained=True)
        self.backbone = nn.Sequential(*list(self.backbone.children())[:-2])
        
        # 空间注意力
        self.spatial_attn = nn.Sequential(
            nn.Conv2d(512, 1, kernel_size=7, padding=3),
            nn.Sigmoid()
        )
        
        # 通道注意力
        self.channel_attn = nn.Sequential(
            nn.AdaptiveAvgPool2d(1),
            nn.Flatten(),
            nn.Linear(512, 128),
            nn.ReLU(),
            nn.Linear(128, 512),
            nn.Sigmoid()
        )
        
        # 输出投影
        self.proj = nn.Sequential(
            nn.AdaptiveAvgPool2d(1),
            nn.Flatten(),
            nn.Linear(512, feature_dim),
            nn.LayerNorm(feature_dim)
        )
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Args:
            x: (B, 3, 64, 64) 眼部图像
        Returns:
            feat: (B, feature_dim) 注意力增强的眼部特征
        """
        feat = self.backbone(x)  # (B, 512, H, W)
        
        # 空间注意力
        spatial_weights = self.spatial_attn(feat)
        feat_spatial = feat * spatial_weights
        
        # 通道注意力
        channel_weights = self.channel_attn(feat)
        feat_channel = feat_spatial * channel_weights.unsqueeze(-1).unsqueeze(-1)
        
        return self.proj(feat_channel)


class CrossModalAttentionFusion(nn.Module):
    """
    跨模态注意力融合 + 可学习门控
    
    三个模态互相查询：
    - 眼部特征查询面部上下文
    - 眼部特征查询头部姿态
    - 面部查询眼部细节
    
    门控决定最终各模态权重（非固定）
    """
    
    def __init__(self, eye_dim, face_dim, head_dim, 
                 hidden_dim=256, num_heads=4):
        super().__init__()
        
        # 投影到统一维度
        self.eye_proj = nn.Linear(eye_dim, hidden_dim)
        self.face_proj = nn.Linear(face_dim, hidden_dim)
        self.head_proj = nn.Linear(head_dim, hidden_dim)
        
        # 跨模态注意力
        self.cross_attn = nn.MultiheadAttention(
            embed_dim=hidden_dim,
            num_heads=num_heads,
            batch_first=True
        )
        
        # 可学习门控
        self.gate = nn.Sequential(
            nn.Linear(hidden_dim * 3, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, 3),  # 三个模态的权重
            nn.Softmax(dim=-1)
        )
    
    def forward(self, eye_feat, face_feat, head_feat):
        """
        Returns:
            fused: (B, hidden_dim) 融合特征
        """
        # 投影
        eye = self.eye_proj(eye_feat)    # (B, hidden)
        face = self.face_proj(face_feat) # (B, hidden)
        head = self.head_proj(head_feat) # (B, hidden)
        
        # 堆叠为序列
        seq = torch.stack([eye, face, head], dim=1)  # (B, 3, hidden)
        
        # 跨模态注意力
        attended, _ = self.cross_attn(seq, seq, seq)
        
        # 门控权重
        concat = torch.cat([eye, face, head], dim=-1)  # (B, hidden*3)
        weights = self.gate(concat)  # (B, 3)
        
        # 加权融合
        fused = (attended * weights.unsqueeze(-1)).sum(dim=1)
        
        return fused


class UncertaintyGazeRegressor(nn.Module):
    """
    不确定性感知视线回归器
    
    使用异方差回归（heteroscedastic regression）：
    - 预测 gaze (θ, φ) 和 variance σ²
    - σ² 大 → 不确定性高 → 预测不可信
    - σ² 小 → 不确定性低 → 预测可信
    
    损失函数：NLL损失
    L = (1/2) * exp(-s) * ||y - ŷ||² + (1/2) * s
    其中 s = log(σ²)
    
    这让模型在难以预测的样本上自动增大σ²
    """
    
    def __init__(self, input_dim, output_dim=2):
        super().__init__()
        
        # 共享特征
        self.shared = nn.Sequential(
            nn.Linear(input_dim, 128),
            nn.ReLU(),
            nn.Dropout(0.2)
        )
        
        # 视线回归头
        self.gaze_head = nn.Linear(128, output_dim)
        
        # 不确定性头（预测log方差）
        self.uncertainty_head = nn.Linear(128, 1)
    
    def forward(self, x):
        feat = self.shared(x)
        gaze = self.gaze_head(feat)
        log_var = self.uncertainty_head(feat)
        uncertainty = torch.exp(log_var)  # σ²
        
        return {
            'gaze': gaze,           # (B, 2) pitch, yaw in radians
            'uncertainty': uncertainty,  # (B, 1) variance
            'confidence': 1.0 / (1.0 + uncertainty)  # (B, 1) confidence score
        }
    
    def loss_function(self, pred_gaze, target_gaze, uncertainty):
        """异方差NLL损失"""
        nll = 0.5 * torch.exp(-torch.log(uncertainty + 1e-8)) * \
              (pred_gaze - target_gaze).pow(2) + \
              0.5 * torch.log(uncertainty + 1e-8)
        return nll.mean()


# 测试
if __name__ == "__main__":
    model = AttentiveGazeModel()
    
    B = 4
    eye = torch.randn(B, 3, 64, 64)
    face = torch.randn(B, 3, 224, 224)
    head = torch.randn(B, 3)
    
    output = model(eye, face, head)
    
    print("=== AttentiveGaze 输出 ===")
    print(f"视线方向: {output['gaze'].shape}")
    print(f"  pitch: {output['gaze'][0, 0]:.4f} rad = {torch.degrees(output['gaze'][0, 0]):.1f}°")
    print(f"  yaw:   {output['gaze'][0, 1]:.4f} rad = {torch.degrees(output['gaze'][0, 1]):.1f}°")
    print(f"不确定性(σ²): {output['uncertainty'][0]:.4f}")
    print(f"置信度: {output['confidence'][0]:.4f}")

class SafeGazeDMS:
    """
    不确定性感知的安全DMS
    
    核心原则：当不确定性高时，不信任预测，采取保守策略
    
    部署策略：
    - σ² < 0.01 → 高置信，直接使用视线预测
    - 0.01 < σ² < 0.05 → 中置信，结合历史数据
    - σ² > 0.05 → 低置信，触发保守策略（假定为分心）
    """
    
    THRESHOLDS = {
        'high_confidence': 0.01,
        'medium_confidence': 0.05,
    }
    
    def __init__(self):
        self.history = []  # 视线历史
        self.window_size = 30  # 1秒@30fps
    
    def process_gaze(self, gaze_result: dict) -> dict:
        """
        处理视线估计结果，做安全决策
        
        Args:
            gaze_result: {'gaze': (2,), 'uncertainty': (1,), 'confidence': (1,)}
        """
        uncertainty = gaze_result['uncertainty'].item()
        gaze = gaze_result['gaze']
        
        if uncertainty < self.THRESHOLDS['high_confidence']:
            # 高置信：直接使用
            action = self.evaluate_attention(gaze)
            return {**action, 'mode': 'high_confidence', 'sigma': uncertainty}
        
        elif uncertainty < self.THRESHOLDS['medium_confidence']:
            # 中置信：结合历史
            if len(self.history) >= 5:
                avg_gaze = torch.stack(self.history[-5:]).mean(dim=0)
                action = self.evaluate_attention(avg_gaze)
                return {**action, 'mode': 'medium_confidence', 'sigma': uncertainty}
            else:
                return self.conservative_action(uncertainty)
        
        else:
            # 低置信：保守策略
            return self.conservative_action(uncertainty)
    
    def evaluate_attention(self, gaze):
        """评估注意力状态"""
        pitch, yaw = gaze[0].item(), gaze[1].item()
        
        # 视线是否在道路区域
        on_road = abs(pitch) < 0.3 and abs(yaw) < 0.4
        
        if not on_road:
            return {'state': 'distracted', 'action': 'warn'}
        else:
            return {'state': 'attentive', 'action': 'none'}
    
    def conservative_action(self, uncertainty):
        """保守策略：假定为分心"""
        return {
            'state': 'unknown',
            'action': 'conservative_warn',
            'mode': 'low_confidence',
            'sigma': uncertainty,
            'reason': f'高不确定性(σ²={uncertainty:.3f})，采取保守策略'
        }


# 测试
if __name__ == "__main__":
    dms = SafeGazeDMS()
    
    # 模拟不同置信度场景
    test_cases = [
        # 高置信，看路
        {'gaze': torch.tensor([0.05, 0.1]), 
         'uncertainty': torch.tensor([0.005]), 
         'confidence': torch.tensor([0.99])},
        # 高置信，看手机
        {'gaze': torch.tensor([0.5, 0.8]),
         'uncertainty': torch.tensor([0.008]),
         'confidence': torch.tensor([0.99])},
        # 低置信（墨镜遮挡）
        {'gaze': torch.tensor([0.1, 0.1]),
         'uncertainty': torch.tensor([0.08]),
         'confidence': torch.tensor([0.92])},
    ]
    
    for i, tc in enumerate(test_cases):
        result = dms.process_gaze(tc)
        print(f"场景{i+1}: {result['mode']} → {result['state']} → {result['action']} (σ²={result['sigma']:.3f})")