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

class GazeAnywhere(nn.Module):
    """
    GazeAnywhere: 可提示视线目标估计模型
    
    论文: CVPR 2026
    架构: 冻结DINOv3 + 冻结文本编码器 + Transformer检测器
    
    核心创新:
    1. 自然语言提示指定目标人物
    2. 端到端，无中间检测阶段
    3. 同时解决定位+在帧判断+视线估计
    """
    
    def __init__(self, 
                 dino_model_name='dinov3_vitl14',
                 text_embed_dim=768,
                 hidden_dim=256,
                 num_queries=100,
                 num_decoder_layers=6):
        super().__init__()
        
        # 冻结视觉编码器 (DINOv3)
        self.image_encoder = self._load_frozen_dino(dino_model_name)
        D_V = self.image_encoder.embed_dim  # 视觉特征维度
        
        # 冻结文本编码器
        self.text_encoder = self._load_frozen_text_encoder()
        D_T = text_embed_dim
        
        # 投影到统一空间
        self.image_proj = nn.Linear(D_V, hidden_dim)
        self.text_proj = nn.Linear(D_T, hidden_dim)
        
        # Transformer检测器（可训练）
        decoder_layer = nn.TransformerDecoderLayer(
            d_model=hidden_dim,
            nhead=8,
            dim_feedforward=hidden_dim * 4,
            dropout=0.1,
            batch_first=True
        )
        self.transformer_decoder = nn.TransformerDecoder(
            decoder_layer, 
            num_layers=num_decoder_layers
        )
        
        # 查询嵌入
        self.queries = nn.Embedding(num_queries, hidden_dim)
        
        # 预测头
        self.bbox_head = MLP(hidden_dim, hidden_dim, 4, 3)  # 头部边界框
        self.presence_head = nn.Linear(hidden_dim, 1)       # 在/不在画面
        self.gaze_head = MLP(hidden_dim, hidden_dim, 2, 3)  # 视线落点(x,y)
        self.gaze_heatmap_head = nn.Linear(hidden_dim, 1)   # 热力图
        
    def _load_frozen_dino(self, name):
        """加载冻结的DINOv3模型"""
        # 实际使用 torchvision 或 hub 加载
        model = torch.hub.load('facebookresearch/dinov3', name)
        for param in model.parameters():
            param.requires_grad = False
        return model
    
    def _load_frozen_text_encoder(self):
        """加载冻结的文本编码器"""
        # 使用 dino.txt 或 CLIP 文本编码器
        model = torch.hub.load('facebookresearch/dinov3', 'dinotxt_base')
        for param in model.parameters():
            param.requires_grad = False
        return model
    
    def forward(self, image: torch.Tensor, 
                text_prompt: str = None,
                visual_prompt: Tuple[float, float] = None) -> dict:
        """
        前向传播
        
        Args:
            image: (B, 3, H, W) 输入图像
            text_prompt: 文本提示字符串
            visual_prompt: (x, y) 归一化坐标提示
            
        Returns:
            dict: {
                'bbox': (B, 4) 头部边界框,
                'presence': (B, 1) 在画面内的概率,
                'gaze_point': (B, 2) 视线落点,
                'gaze_heatmap': (B, H//14, W//14) 视线热力图
            }
        """
        B = image.shape[0]
        
        # 1. 提取冻结视觉特征
        with torch.no_grad():
            visual_features = self.image_encoder(image)
            # DINOv3输出: [CLS] + patch tokens
            # visual_features: (B, 1+N_V, D_V)
        
        visual_features = self.image_proj(visual_features)
        
        # 2. 提取文本特征
        if text_prompt is not None:
            with torch.no_grad():
                text_features = self.text_encoder(text_prompt)
            text_features = self.text_proj(text_features)
            # (B, D_T) -> (B, 1, hidden_dim)
            text_features = text_features.unsqueeze(1)
        elif visual_prompt is not None:
            # 将坐标提示映射到embedding
            coord_embed = self._coord_to_embed(visual_prompt, visual_features, image)
            text_features = coord_embed  # (B, 1, hidden_dim)
        
        # 3. Transformer解码
        queries = self.queries.weight.unsqueeze(0).expand(B, -1, -1)
        
        # 将文本特征作为key，视觉特征+文本作为memory
        memory = torch.cat([visual_features, text_features], dim=1)
        
        decoded = self.transformer_decoder(queries, memory)
        # decoded: (B, num_queries, hidden_dim)
        
        # 取第一个query的输出作为主体特征
        subject_feat = decoded[:, 0, :]  # (B, hidden_dim)
        
        # 4. 预测
        bbox = self.bbox_head(subject_feat).sigmoid()  # 归一化坐标
        presence = self.presence_head(subject_feat).sigmoid()
        gaze_point = self.gaze_head(subject_feat).sigmoid()
        
        # 热力图（简化版，实际用所有query的attention）
        heatmap_logits = self.gaze_heatmap_head(decoded)  # (B, num_queries, 1)
        gaze_heatmap = heatmap_logits.sigmoid()
        
        return {
            'bbox': bbox,
            'presence': presence,
            'gaze_point': gaze_point,
            'gaze_heatmap': gaze_heatmap
        }
    
    def _coord_to_embed(self, coord, visual_features, image):
        """将坐标提示转换为embedding（通过位置编码）"""
        B, N, D = visual_features.shape
        H_patches = image.shape[2] // 14
        W_patches = image.shape[3] // 14
        
        x_idx = int(coord[0] * W_patches)
        y_idx = int(coord[1] * H_patches)
        token_idx = y_idx * W_patches + x_idx + 1  # +1 for CLS
        
        coord_embed = visual_features[:, token_idx:token_idx+1, :]
        return coord_embed


class MLP(nn.Module):
    """多层感知机"""
    def __init__(self, input_dim, hidden_dim, output_dim, num_layers):
        super().__init__()
        h = [hidden_dim] * (num_layers - 1)
        self.layers = nn.ModuleList(
            nn.Linear(n, k) for n, k in zip([input_dim] + h, h + [output_dim])
        )
    
    def forward(self, x):
        for i, layer in enumerate(self.layers):
            x = torch.relu(layer(x)) if i < len(self.layers) - 1 else layer(x)
        return x


# 测试代码
if __name__ == "__main__":
    model = GazeAnywhere()
    model.eval()
    
    # 模拟输入
    image = torch.randn(2, 3, 224, 224)
    text_prompt = "the person wearing a red shirt"
    
    with torch.no_grad():
        output = model(image, text_prompt=text_prompt)
    
    print("=== GazeAnywhere 测试 ===")
    print(f"输入图像: {image.shape}")
    print(f"文本提示: '{text_prompt}'")
    print(f"预测bbox: {output['bbox'].shape} → {output['bbox'][0].tolist()}")
    print(f"在画面内: {output['presence'][0].item():.4f}")
    print(f"视线落点: {output['gaze_point'][0].tolist()}")
    print(f"视线热力图: {output['gaze_heatmap'].shape}")
    print()
    print("论文核心价值: 端到端，无中间检测阶段")
    print("→ IMS应用: 可用自然语言指定驾驶员，无需人脸检测前置")

class DMSGazeIntegration:
    """
    DMS + GazeAnywhere集成方案
    
    用途: 驾驶员视线落点检测
    优势: 无需人脸检测前置，支持墨镜/口罩场景
    """
    
    # 视线区域定义（归一化坐标）
    GAZE_ZONES = {
        'road_ahead':    [(0.3, 0.4), (0.7, 0.6)],   # 前方道路
        'left_mirror':   [(0.0, 0.3), (0.2, 0.5)],   # 左后视镜
        'right_mirror':  [(0.8, 0.3), (1.0, 0.5)],   # 右后视镜
        'dashboard':     [(0.3, 0.6), (0.7, 1.0)],   # 仪表盘
        'infotainment':  [(0.4, 0.7), (0.6, 0.9)],   # 中控屏
        'phone':         [(0.2, 0.6), (0.4, 0.8)],   # 手机区域
        'passenger':     [(0.6, 0.2), (0.9, 0.5)]    # 乘客区域
    }
    
    def __init__(self):
        self.model = GazeAnywhere()
        self.distraction_timer = 0
        
    def detect_distraction(self, image, gaze_point):
        """检测分心状态"""
        x, y = gaze_point
        zone = self._classify_zone(x, y)
        
        if zone in ['left_mirror', 'right_mirror', 'dashboard']:
            # 正常区域，重置计时
            self.distraction_timer = 0
            return {'zone': zone, 'distraction': False}
        elif zone in ['phone', 'passenger', 'infotainment']:
            # 分心区域
            self.distraction_timer += 1
            if self.distraction_timer > 90:  # 3秒@30fps
                return {
                    'zone': zone,
                    'distraction': True,
                    'duration_s': self.distraction_timer / 30,
                    'level': 'warning' if self.distraction_timer < 180 else 'danger'
                }
        
        return {'zone': zone, 'distraction': False}