"""
航空眼动追踪培训方法 → IMS 驾驶员监测的跨域映射

航空培训关注点：
1. 仪表扫描效率（是否看了该看的仪表）
2. 注意力分配（任务优先级）
3. 情境意识评估
4. 疲劳/过载检测

IMS 可迁移的方法：
1. 视线区域映射 → 驾驶场景注意力评估
2. 扫描模式分析 → 驾驶员注意力分配评分
3. 情境意识指标 → 功能就绪度评估
4. 疲劳检测 → 直接复用
"""

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


class GazeRegion(IntEnum):
    """视线区域（可同时用于航空和汽车）"""
    # 航空仪表区域
    PRIMARY_FLIGHT_DISPLAY = 0    # 主飞行显示
    NAVIGATION_DISPLAY = 1        # 导航显示
    ENGINE_INSTRUMENTS = 2         # 发动机仪表
    OVERHEAD_PANEL = 3             # 顶板
    # 汽车区域映射
    ROAD = 10                      # 道路
    REARVIEW_MIRROR = 11           # 后视镜
    INSTRUMENT_CLUSTER = 12       # 仪表盘
    INFOTAINMENT = 13             # 中控屏
    STEERING_WHEEL = 14           # 方向盘


@dataclass
class GazeSample:
    """眼动采样点"""
    timestamp: float  # 秒
    pitch: float      # 俯仰角 (度)
    yaw: float        # 偏航角 (度)
    region: GazeRegion = None  # 映射的视线区域


class ScanPatternAnalyzer:
    """
    扫描模式分析器
    
    从航空培训方法迁移到 IMS 驾驶员监测
    
    核心指标：
    1. 区域停留时间分布（注意力分配）
    2. 扫描频率（仪表/区域切换速率）
    3. 任务相关注视比例
    4. 情境意识评分
    """
    
    def __init__(self, window_sec=60.0, fps=30):
        self.window_sec = window_sec
        self.fps = fps
        self.window_samples = int(window_sec * fps)
    
    def analyze_attention_distribution(self, 
                                        samples: List[GazeSample]) -> dict:
        """
        分析注意力分配
        
        Args:
            samples: 眼动采样序列
            
        Returns:
            各区域停留时间比例 + 注意力分散度
        """
        if len(samples) < 2:
            return {}
        
        # 计算每个区域停留时间
        region_time = {}
        total_time = 0
        
        for i in range(1, len(samples)):
            dt = samples[i].timestamp - samples[i-1].timestamp
            region = samples[i-1].region
            if region is not None:
                region_time[region] = region_time.get(region, 0) + dt
                total_time += dt
        
        # 转为百分比
        attention_dist = {
            r: t/total_time * 100 
            for r, t in region_time.items()
        }
        
        # 注意力分散度（Shannon 熵）
        probs = np.array(list(attention_dist.values())) / 100
        probs = probs[probs > 0]
        entropy = -np.sum(probs * np.log2(probs))
        max_entropy = np.log2(len(GazeRegion))
        normalized_entropy = entropy / max_entropy if max_entropy > 0 else 0
        
        return {
            'attention_distribution': attention_dist,
            'attention_entropy': normalized_entropy,  # 0=高度集中, 1=高度分散
            'total_regions_visited': len(attention_dist)
        }
    
    def compute_scan_rate(self, samples: List[GazeSample]) -> float:
        """
        计算扫描频率（区域切换/秒）
        
        航空参考值：
        - 熟练飞行员：0.3-0.5 Hz（结构化扫描）
        - 疲劳飞行员：0.1-0.2 Hz（凝视增多）
        
        汽车参考值：
        - 正常驾驶：0.5-1.0 Hz（道路+后视镜+仪表）
        - 疲劳驾驶：0.1-0.3 Hz（凝视道路）
        """
        transitions = 0
        for i in range(1, len(samples)):
            if samples[i].region != samples[i-1].region:
                transitions += 1
        
        duration = samples[-1].timestamp - samples[0].timestamp
        return transitions / duration if duration > 0 else 0
    
    def compute_situational_awareness(self, 
                                       samples: List[GazeSample],
                                       critical_regions: List[GazeRegion]) -> float:
        """
        情境意识评分
        
        关键区域注视比例 × 扫描频率归一化
        
        Args:
            samples: 眼动采样
            critical_regions: 关键区域列表（如道路、后视镜）
            
        Returns:
            SA_score: 情境意识评分 (0-100)
        """
        dist = self.analyze_attention_distribution(samples)
        
        # 关键区域注视比例
        critical_time = sum(
            dist['attention_distribution'].get(r, 0)
            for r in critical_regions
        )
        critical_ratio = critical_time / 100  # 0-1
        
        # 扫描频率
        scan_rate = self.compute_scan_rate(samples)
        scan_norm = min(scan_rate / 1.0, 1.0)  # 归一化到 0-1
        
        # SA 评分
        sa_score = (critical_ratio * 0.7 + scan_norm * 0.3) * 100
        
        return sa_score


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

if __name__ == "__main__":
    np.random.seed(42)
    analyzer = ScanPatternAnalyzer(window_sec=60, fps=30)
    
    # 模拟正常驾驶（频繁扫描道路+后视镜）
    normal_samples = []
    for i in range(1800):  # 60秒 × 30fps
        t = i / 30
        r = np.random.choice(
            [GazeRegion.ROAD, GazeRegion.REARVIEW_MIRROR,
             GazeRegion.INSTRUMENT_CLUSTER],
            p=[0.7, 0.15, 0.15]
        )
        normal_samples.append(GazeSample(
            timestamp=t,
            pitch=np.random.normal(-5, 3),
            yaw=np.random.normal(0, 10),
            region=r
        ))
    
    # 模拟疲劳驾驶（凝视道路，少扫描）
    fatigue_samples = []
    for i in range(1800):
        t = i / 30
        r = GazeRegion.ROAD if np.random.random() > 0.05 else GazeRegion.INSTRUMENT_CLUSTER
        fatigue_samples.append(GazeSample(
            timestamp=t,
            pitch=np.random.normal(-3, 1),
            yaw=np.random.normal(0, 3),
            region=r
        ))
    
    # 分析
    critical = [GazeRegion.ROAD, GazeRegion.REARVIEW_MIRROR]
    
    normal_result = analyzer.analyze_attention_distribution(normal_samples)
    normal_scan = analyzer.compute_scan_rate(normal_samples)
    normal_sa = analyzer.compute_situational_awareness(normal_samples, critical)
    
    fatigue_result = analyzer.analyze_attention_distribution(fatigue_samples)
    fatigue_scan = analyzer.compute_scan_rate(fatigue_samples)
    fatigue_sa = analyzer.compute_situational_awareness(fatigue_samples, critical)
    
    print("=== 正常驾驶 ===")
    print(f"  注意力分配: {normal_result['attention_distribution']}")
    print(f"  注意力熵: {normal_result['attention_entropy']:.3f}")
    print(f"  扫描频率: {normal_scan:.2f} Hz")
    print(f"  情境意识评分: {normal_sa:.1f}/100")
    
    print("\n=== 疲劳驾驶 ===")
    print(f"  注意力分配: {fatigue_result['attention_distribution']}")
    print(f"  注意力熵: {fatigue_result['attention_entropy']:.3f}")
    print(f"  扫描频率: {fatigue_scan:.2f} Hz")
    print(f"  情境意识评分: {fatigue_sa:.1f}/100")