import torch
import torch.nn as nn

class SGAPGaze(nn.Module):
    """
    SGAP-Gaze: Scene Grid Attention Point-of-Gaze Estimation
    
    创新: 将场景图像分成网格，用注意力机制
    学习驾驶员注视落在哪个网格区域
    
    输入: 驾驶员面部图 + 场景图
    输出: 注视点坐标(x, y) + 网格注意力图
    """
    
    def __init__(self, grid_size: int = 8, feat_dim: int = 256):
        super().__init__()
        self.grid_size = grid_size
        
        # 驾驶员面部编码器
        self.face_encoder = nn.Sequential(
            nn.Conv2d(3, 32, 3, stride=2, padding=1), nn.ReLU(),
            nn.Conv2d(32, 64, 3, stride=2, padding=1), nn.ReLU(),
            nn.Conv2d(64, 128, 3, stride=2, padding=1), nn.ReLU(),
            nn.AdaptiveAvgPool2d((1, 1)), nn.Flatten(),
            nn.Linear(128, feat_dim)
        )
        
        # 场景编码器（网格化）
        self.scene_encoder = nn.Sequential(
            nn.Conv2d(3, 32, 3, stride=2, padding=1), nn.ReLU(),
            nn.Conv2d(32, 64, 3, stride=2, padding=1), nn.ReLU(),
            nn.Conv2d(64, 128, 3, stride=2, padding=1), nn.ReLU(),
            nn.Conv2d(128, 256, 3, padding=1), nn.ReLU(),
        )
        
        # 场景网格注意力
        self.grid_attention = nn.MultiheadAttention(
            embed_dim=feat_dim,
            num_heads=8,
            batch_first=True
        )
        
        # 注视点回归头
        self.pog_head = nn.Sequential(
            nn.Linear(feat_dim * 2, 128), nn.ReLU(),
            nn.Linear(128, 2)  # (x, y) 归一化坐标
        )
        
    def forward(self, face_img: torch.Tensor, scene_img: torch.Tensor) -> dict:
        """
        Args:
            face_img: (B, 3, 128, 128) 驾驶员面部
            scene_img: (B, 3, 256, 256) 前方场景
        
        Returns:
            pog: (B, 2) 注视点坐标
            attention_map: (B, grid*grid) 网格注意力
        """
        # 驾驶员特征
        face_feat = self.face_encoder(face_img)  # (B, 256)
        
        # 场景网格特征
        scene_feat = self.scene_encoder(scene_img)  # (B, 256, H', W')
        B, C, H, W = scene_feat.shape
        
        # 展平为网格序列
        grid_feats = scene_feat.view(B, C, H * W).permute(0, 2, 1)  # (B, H*W, 256)
        
        # 用面部特征查询场景网格
        query = face_feat.unsqueeze(1)  # (B, 1, 256)
        attn_out, attn_weights = self.grid_attention(
            query, grid_feats, grid_feats
        )
        
        # 注视点回归
        combined = torch.cat([face_feat, attn_out.squeeze(1)], dim=1)
        pog = self.pog_head(combined)
        
        return {
            'pog': pog,
            'attention_map': attn_weights.squeeze(1),
            'grid_size': (H, W)
        }


# 测试
if __name__ == "__main__":
    model = SGAPGaze(grid_size=8, feat_dim=256)
    
    face = torch.randn(4, 3, 128, 128)
    scene = torch.randn(4, 3, 256, 256)
    
    out = model(face, scene)
    
    print(f"面部输入: {face.shape}")
    print(f"场景输入: {scene.shape}")
    print(f"注视点: {out['pog'].shape} (归一化坐标)")
    print(f"注意力图: {out['attention_map'].shape}")
    print(f"网格尺寸: {out['grid_size']}")

DRIVING_GAZE_ZONES = {
    'road_ahead':    {'bbox': (0.3, 0.2, 0.7, 0.6), 'risk': 'safe'},
    'left_mirror':   {'bbox': (0.0, 0.3, 0.15, 0.5), 'risk': 'safe'},
    'right_mirror':  {'bbox': (0.85, 0.3, 1.0, 0.5), 'risk': 'safe'},
    'instrument':    {'bbox': (0.3, 0.6, 0.7, 0.8), 'risk': 'caution'},
    'center_screen': {'bbox': (0.3, 0.5, 0.7, 0.7), 'risk': 'distraction'},
    'phone':         {'bbox': (0.0, 0.5, 0.3, 0.8), 'risk': 'danger'},
    'passenger':     {'bbox': (0.0, 0.0, 0.3, 0.5), 'risk': 'moderate'},
}

def classify_gaze_zone(pog_x: float, pog_y: float) -> str:
    """根据注视点分类驾驶区域"""
    for zone, info in DRIVING_GAZE_ZONES.items():
        x1, y1, x2, y2 = info['bbox']
        if x1 <= pog_x <= x2 and y1 <= pog_y <= y2:
            return zone
    return 'other'