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

@dataclass
class GazeFeature:
    """眼动特征数据结构"""
    gaze_x: float          # 注视点 X 坐标 (归一化)
    gaze_y: float          # 注视点 Y 坐标 (归一化)
    pupil_diameter: float  # 瞳孔直径 (mm)
    blink_rate: float      # 眨眼频率 (次/分钟)
    perclos: float         # PERCLOS 值 (0-1)
    fixation_stability: float  # 注视稳定性指标


@dataclass
class FacialFeature:
    """面部行为特征数据结构"""
    au_intensity: np.ndarray   # Facial Action Unit 强度 (17维)
    head_pitch: float          # 头部俯仰角 (度)
    head_yaw: float           # 头部偏航角 (度)
    head_roll: float          # 头部翻滚角 (度)
    mouth_openness: float     # 嘴巴开度
    eye_openness_l: float     # 左眼开度 (0-1)
    eye_openness_r: float     # 右眼开度 (0-1)


def compute_impairment_score(
    gaze_features: List[GazeFeature],
    facial_features: List[FacialFeature],
    window_sec: int = 30,
    fps: int = 30
) -> float:
    """
    计算酒精损伤综合评分
    
    基于 Smart Eye AIS 系统的检测原理复现：
    - 眼动稳定性分析
    - 面部行为模式分析  
    - 头部姿态晃动分析
    - PERCLOS 异常检测
    
    Args:
        gaze_features: 眼动特征序列
        facial_features: 面部特征序列
        window_sec: 滑动窗口秒数
        fps: 帧率
    
    Returns:
        impairment_score: 损伤评分 (0-1), >0.7 判定为疑似损伤
    """
    window_size = int(window_sec * fps)
    
    if len(gaze_features) < window_size:
        return 0.0
    
    scores = []
    for i in range(len(gaze_features) - window_size):
        gaze_window = gaze_features[i:i+window_size]
        facial_window = facial_features[i:i+window_size]
        
        # 1. 注视稳定性评分 (酒精损伤导致注视漂移)
        gaze_x = np.array([g.gaze_x for g in gaze_window])
        gaze_y = np.array([g.gaze_y for g in gaze_window])
        gaze_stability = 1.0 - min(np.std(gaze_x) + np.std(gaze_y), 1.0)
        
        # 2. PERCLOS 异常评分
        perclos_vals = np.array([g.perclos for g in gaze_window])
        avg_perclos = np.mean(perclos_vals)
        # 正常 PERCLOS < 0.15, 酒精损伤常 > 0.25
        perclos_score = max(0, (avg_perclos - 0.15) / 0.35)
        
        # 3. 眨眼频率异常评分
        blink_rates = np.array([g.blink_rate for g in gaze_window])
        avg_blink = np.mean(blink_rates)
        # 正常眨眼 15-20次/分, 酒精损伤可能 25-40次/分
        blink_score = max(0, (avg_blink - 20) / 20)
        
        # 4. 头部晃动评分
        pitch = np.array([f.head_pitch for f in facial_window])
        yaw = np.array([f.head_yaw for f in facial_window])
        roll = np.array([f.head_roll for f in facial_window])
        # 计算头部运动的高频成分
        pitch_var = np.std(np.diff(pitch))
        yaw_var = np.std(np.diff(yaw))
        roll_var = np.std(np.diff(roll))
        head_instability = min((pitch_var + yaw_var + roll_var) / 10, 1.0)
        
        # 5. 面部表情松弛评分 (AU 分析)
        au_intensities = np.array([f.au_intensity for f in facial_window])
        # AU1(内眉上提), AU4(眉头下压), AU15(嘴角下压) 
        # 酒精损伤时面部 AU 整体强度降低
        avg_au_intensity = np.mean(au_intensities, axis=0)
        facial_relaxation = 1.0 - np.mean(avg_au_intensity)
        
        # 加权综合评分
        total_score = (
            0.25 * (1.0 - gaze_stability) +     # 注视不稳定
            0.25 * perclos_score +               # PERCLOS 异常
            0.15 * blink_score +                 # 眨眼异常
            0.20 * head_instability +            # 头部晃动
            0.15 * facial_relaxation             # 面部松弛
        )
        scores.append(total_score)
    
    return float(np.mean(scores[-10:]))  # 取最近10个窗口均值


# ============ 测试用例 ============
if __name__ == "__main__":
    np.random.seed(42)
    
    # 模拟正常驾驶 (30秒, 30fps = 900帧)
    normal_gaze = [GazeFeature(
        gaze_x=np.random.normal(0.5, 0.02),
        gaze_y=np.random.normal(0.5, 0.02),
        pupil_diameter=4.0,
        blink_rate=18,
        perclos=0.08,
        fixation_stability=0.95
    ) for _ in range(900)]
    
    normal_facial = [FacialFeature(
        au_intensity=np.random.uniform(0.3, 0.7, 17),
        head_pitch=np.random.normal(0, 0.5),
        head_yaw=np.random.normal(0, 0.5),
        head_roll=np.random.normal(0, 0.3),
        mouth_openness=0.1,
        eye_openness_l=0.8,
        eye_openness_r=0.8
    ) for _ in range(900)]
    
    # 模拟酒精损伤驾驶
    drunk_gaze = [GazeFeature(
        gaze_x=np.random.normal(0.5, 0.08),  # 注视漂移增大
        gaze_y=np.random.normal(0.5, 0.08),
        pupil_diameter=4.5,
        blink_rate=32,                         # 眨眼频率升高
        perclos=0.28,                          # PERCLOS 升高
        fixation_stability=0.65
    ) for _ in range(900)]
    
    drunk_facial = [FacialFeature(
        au_intensity=np.random.uniform(0.1, 0.3, 17),  # AU 强度降低
        head_pitch=np.random.normal(0, 2.0),   # 头部晃动增大
        head_yaw=np.random.normal(0, 2.0),
        head_roll=np.random.normal(0, 1.5),
        mouth_openness=0.05,
        eye_openness_l=0.65,
        eye_openness_r=0.65
    ) for _ in range(900)]
    
    # 测试
    normal_score = compute_impairment_score(normal_gaze, normal_facial)
    drunk_score = compute_impairment_score(drunk_gaze, drunk_facial)
    
    print(f"正常驾驶损伤评分: {normal_score:.3f} (阈值: 0.7)")
    print(f"酒精损伤评分: {drunk_score:.3f} (阈值: 0.7)")
    print(f"判定结果 - 正常: {'通过' if normal_score < 0.7 else '疑似损伤'}")
    print(f"判定结果 - 酒精: {'疑似损伤' if drunk_score >= 0.7 else '通过'}")

class ImpairmentDetectionModule:
    """
    IMS 损伤检测模块 - 可插入现有 DMS 管线
    
    集成方式:
    1. 复用现有眼动追踪模块的输出
    2. 复用现有面部关键点检测的输出  
    3. 新增损伤评分模块
    4. 新增损伤告警逻辑
    """
    
    def __init__(self, config: dict):
        self.window_sec = config.get('window_sec', 30)
        self.threshold = config.get('threshold', 0.7)
        self.fps = config.get('fps', 30)
        
        # 特征缓冲区
        self.gaze_buffer = []
        self.facial_buffer = []
        self.max_buffer = self.window_sec * self.fps
        
        # 状态
        self.alert_level = 0  # 0=正常, 1=一级, 2=二级, 3=三级
        self.alert_duration = 0  # 损伤持续秒数
        
    def update(self, gaze_data, facial_data) -> int:
        """
        每帧调用 - 更新损伤检测状态
        
        Args:
            gaze_data: 来自眼动追踪模块的输出
            facial_data: 来自面部关键点模块的输出
            
        Returns:
            alert_level: 当前告警等级
        """
        self.gaze_buffer.append(gaze_data)
        self.facial_buffer.append(facial_data)
        
        if len(self.gaze_buffer) > self.max_buffer:
            self.gaze_buffer.pop(0)
            self.facial_buffer.pop(0)
        
        if len(self.gaze_buffer) < self.max_buffer:
            return 0
        
        score = compute_impairment_score(
            self.gaze_buffer, 
            self.facial_buffer,
            self.window_sec,
            self.fps
        )
        
        if score >= self.threshold:
            self.alert_duration += 1 / self.fps
            if self.alert_duration > 20:
                self.alert_level = 3  # 三级：安全停车
            elif self.alert_duration > 15:
                self.alert_level = 2  # 二级：持续+触觉
            elif self.alert_duration > 10:
                self.alert_level = 1  # 一级：视觉+听觉
        else:
            self.alert_duration = max(0, self.alert_duration - 2/self.fps)
            if self.alert_duration == 0:
                self.alert_level = 0
        
        return self.alert_level