import numpy as np
from typing import Tuple, List, Optional
from dataclasses import dataclass
from enum import IntEnum

class AlertLevel(IntEnum):
    NORMAL = 0
    VISUAL_HINT = 1
    AR_GUIDE = 2
    AUDIO_ALERT = 3
    HAPTIC_ALERT = 4

@dataclass
class GazePoint:
    x: float  # 归一化 0-1
    y: float  # 归一化 0-1
    timestamp: float
    confidence: float

@dataclass
class RiskZone:
    """道路风险区域"""
    x: float       # 风险区域中心 X
    y: float       # 风险区域中心 Y
    radius: float  # 区域半径
    risk_type: str # risk type
    severity: int  # 1-3
    distance_m: float  # 距离（米）

class ARHUDAttentionGuide:
    """
    AR-HUD 注意力引导系统
    
    核心逻辑：
    1. DMS 实时追踪注视点
    2. 检测注意力偏离风险区域
    3. 通过 AR-HUD 引导注意力回归
    4. 闭环验证引导效果
    
    基于 Frontiers 2025 AR-HUD 研究和 
    Aviation Today 2026 下一代 HUD 报告
    """
    
    def __init__(self, config: dict = None):
        config = config or {}
        self.gaze_history: List[GazePoint] = []
        self.max_history = config.get('max_history', 150)  # 5秒@30fps
        
        # 注意力偏离阈值
        self.deviation_threshold = config.get('deviation_threshold', 0.15)
        self.duration_threshold = config.get('duration_threshold', 2.0)  # 秒
        
        # 引导效果
        self.current_guide_level = AlertLevel.NORMAL
        self.guide_start_time = None
        
    def update(
        self,
        gaze: GazePoint,
        risk_zones: List[RiskZone]
    ) -> dict:
        """
        每帧更新
        
        Args:
            gaze: 当前注视点
            risk_zones: 当前道路风险区域列表
            
        Returns:
            guide_command: AR-HUD 引导指令
        """
        self.gaze_history.append(gaze)
        if len(self.gaze_history) > self.max_history:
            self.gaze_history.pop(0)
        
        # 1. 检查注视是否偏离风险区域
        attention_status = self._check_attention(gaze, risk_zones)
        
        # 2. 确定引导等级
        if attention_status['deviated']:
            if self.guide_start_time is None:
                self.guide_start_time = gaze.timestamp
            
            duration = gaze.timestamp - self.guide_start_time
            
            if duration > self.duration_threshold * 3:
                level = AlertLevel.HAPTIC_ALERT
            elif duration > self.duration_threshold * 2:
                level = AlertLevel.AUDIO_ALERT
            elif duration > self.duration_threshold:
                level = AlertLevel.AR_GUIDE
            else:
                level = AlertLevel.VISUAL_HINT
        else:
            self.guide_start_time = None
            level = AlertLevel.NORMAL
        
        self.current_guide_level = level
        
        # 3. 生成 AR-HUD 渲染指令
        render_cmd = self._generate_render_command(
            level, gaze, risk_zones, attention_status
        )
        
        return {
            'alert_level': int(level),
            'attention_status': attention_status,
            'render_command': render_cmd
        }
    
    def _check_attention(
        self,
        gaze: GazePoint,
        risk_zones: List[RiskZone]
    ) -> dict:
        """检查注意力是否指向风险区域"""
        if not risk_zones:
            return {'deviated': False, 'nearest_risk_dist': float('inf')}
        
        # 计算注视点到各风险区域的距离
        distances = []
        for zone in risk_zones:
            dist = np.sqrt(
                (gaze.x - zone.x)**2 + (gaze.y - zone.y)**2
            )
            distances.append((dist, zone))
        
        distances.sort(key=lambda x: x[0])
        nearest_dist, nearest_zone = distances[0]
        
        # 判定是否偏离
        deviated = nearest_dist > self.deviation_threshold
        
        return {
            'deviated': deviated,
            'nearest_risk_dist': nearest_dist,
            'nearest_risk_type': nearest_zone.risk_type,
            'nearest_risk_severity': nearest_zone.severity,
            'nearest_risk_distance_m': nearest_zone.distance_m
        }
    
    def _generate_render_command(
        self,
        level: AlertLevel,
        gaze: GazePoint,
        risk_zones: List[RiskZone],
        attention_status: dict
    ) -> dict:
        """生成 AR-HUD 渲染指令"""
        cmd = {
            'mode': 'normal',
            'elements': []
        }
        
        if level == AlertLevel.NORMAL:
            cmd['mode'] = 'standard_info'
            
        elif level == AlertLevel.VISUAL_HINT:
            # 轻微视觉提示：风险区域边框闪烁
            for zone in risk_zones:
                if zone.severity >= 2:
                    cmd['elements'].append({
                        'type': 'border_highlight',
                        'x': zone.x, 'y': zone.y,
                        'radius': zone.radius,
                        'color': 'yellow',
                        'blink': True,
                        'blink_rate': 2  # Hz
                    })
            cmd['mode'] = 'hint'
            
        elif level == AlertLevel.AR_GUIDE:
            # AR 箭头引导：指向风险区域
            for zone in risk_zones:
                if zone.severity >= 2:
                    cmd['elements'].append({
                        'type': 'arrow',
                        'from': {'x': gaze.x, 'y': gaze.y},
                        'to': {'x': zone.x, 'y': zone.y},
                        'color': 'orange',
                        'animated': True,
                        'label': zone.risk_type
                    })
            cmd['mode'] = 'guidance'
            
        elif level == AlertLevel.AUDIO_ALERT:
            # 音频告警 + AR 高亮
            for zone in risk_zones:
                cmd['elements'].append({
                    'type': 'full_highlight',
                    'x': zone.x, 'y': zone.y,
                    'radius': zone.radius,
                    'color': 'red',
                    'pulse': True
                })
            cmd['mode'] = 'warning'
            cmd['audio'] = True
            
        elif level == AlertLevel.HAPTIC_ALERT:
            # 触觉告警 + 全屏警告
            cmd['mode'] = 'critical'
            cmd['audio'] = True
            cmd['haptic'] = True
            cmd['elements'].append({
                'type': 'fullscreen_flash',
                'color': 'red'
            })
        
        return cmd


# ============ 测试 ============

if __name__ == "__main__":
    guide = ARHUDAttentionGuide()
    
    print("=== AR-HUD 注意力引导系统 ===")
    print(f"偏离阈值: {guide.deviation_threshold}")
    print(f"持续时间阈值: {guide.duration_threshold}s")
    
    # 模拟场景
    risk_zones = [
        RiskZone(0.5, 0.7, 0.1, 'pedestrian', 3, 30),
        RiskZone(0.8, 0.5, 0.08, 'vehicle_cut_in', 2, 50),
    ]
    
    # 场景1: 正常注视前方
    gaze_normal = GazePoint(0.5, 0.65, 0.0, 0.95)
    result1 = guide.update(gaze_normal, risk_zones)
    print(f"\n场景1 (正常注视):")
    print(f"  告警等级: {result1['alert_level']}")
    print(f"  偏离: {result1['attention_status']['deviated']}")
    print(f"  模式: {result1['render_command']['mode']}")
    
    # 场景2: 注视偏离（看手机）
    guide.guide_start_time = None
    for i in range(60):  # 2秒
        gaze_distracted = GazePoint(0.2, 0.3, i/30, 0.9)
        result2 = guide.update(gaze_distracted, risk_zones)
    
    print(f"\n场景2 (持续偏离2秒):")
    print(f"  告警等级: {result2['alert_level']}")
    print(f"  模式: {result2['render_command']['mode']}")
    print(f"  渲染元素数: {len(result2['render_command']['elements'])}")
    
    # 场景3: 持续偏离5秒
    guide.guide_start_time = None
    for i in range(150):  # 5秒
        gaze_long = GazePoint(0.15, 0.25, i/30, 0.85)
        result3 = guide.update(gaze_long, risk_zones)
    
    print(f"\n场景3 (持续偏离5秒):")
    print(f"  告警等级: {result3['alert_level']}")
    print(f"  模式: {result3['render_command']['mode']}")
    print(f"  音频告警: {result3['render_command'].get('audio', False)}")
    print(f"  触觉告警: {result3['render_command'].get('haptic', False)}")
    
    print(f"\n=== 告警等级说明 ===")
    for level in AlertLevel:
        print(f"  Level {level.value}: {level.name}")