import numpy as np
from dataclasses import dataclass
from typing import Optional

@dataclass
class FacialCues:
    """面部视觉特征"""
    perclos: float  # 眼睑闭合百分比
    head_sway_amplitude: float  # 头部摇晃幅度
    head_sway_frequency: float  # 摇晃频率
    facial_expression_score: float  # 表情活跃度
    pupil_response_latency: Optional[float]  # 瞳孔反应延迟(ms)


@dataclass
class BehavioralCues:
    """驾驶行为特征"""
    steering_reversal_rate: float  # 方向盘修正频率
    lane_deviation_std: float  # 车道偏离标准差
    speed_variation: float  # 速度波动
    brake_reaction_time: float  # 刹车反应时间


@dataclass 
class SobrietyAssessment:
    """酒驾损伤评估结果"""
    impairment_score: float  # 0-1, 1=严重损伤
    bac_estimate: float  # BAC 估计值
    confidence: float
    primary_indicators: list  # 主要贡献指标
    recommendation: str  # 'safe' | 'warning' | 'unsafe'


class VisionBasedSobrietyDetector:
    """
    纯视觉酒驾损伤检测器
    
    基于 Greer et al. (2023) 的研究发现：
    - 面部特征 + 驾驶行为 融合检测
    - BAC 阈值判断：0.02% 安全, 0.05% 警告, 0.08% 危险
    """
    
    # BAC 阈值
    BAC_SAFE = 0.02
    BAC_WARNING = 0.05
    BAC_UNSAFE = 0.08
    
    def __init__(self):
        # 各特征的权重（基于论文检测精度）
        self.facial_weights = {
            'perclos': 0.30,      # 检测精度 85%
            'head_sway': 0.25,    # 检测精度 80%
            'expression': 0.10,   # 检测精度 65%
            'pupil_latency': 0.15, # 检测精度 75%
        }
        self.behavioral_weights = {
            'steering': 0.25,     # 检测精度 78%
            'lane_dev': 0.30,     # 检测精度 82%
            'speed_var': 0.15,    # 检测精度 70%
            'brake_rt': 0.20,     # 检测精度 75%
        }
        
        # 基线特征（清醒状态）
        self.baseline = {
            'perclos': 0.08,           # 正常 PERCLOS ~8%
            'head_sway_amp': 2.0,       # mm
            'head_sway_freq': 0.3,      # Hz
            'expression_score': 0.6,    # 0-1
            'pupil_latency': 250,       # ms
            'steering_rate': 30,       # corrections/min
            'lane_dev_std': 0.15,       # m
            'speed_var': 2.0,           # km/h std
            'brake_rt': 0.8,           # s
        }
    
    def assess(self, facial: FacialCues, 
               behavioral: BehavioralCues) -> SobrietyAssessment:
        """
        评估酒精损伤程度
        
        Args:
            facial: 面部视觉特征
            behavioral: 驾驶行为特征
            
        Returns:
            assessment: 损伤评估结果
        """
        # 计算各特征相对基线的偏离程度
        deviations = {}
        
        # 面部特征偏离
        deviations['perclos'] = self._deviation(
            facial.perclos, self.baseline['perclos'], direction='increase'
        )
        deviations['head_sway'] = self._deviation(
            facial.head_sway_amplitude, self.baseline['head_sway_amp'], 
            direction='increase'
        )
        deviations['expression'] = self._deviation(
            facial.facial_expression_score, self.baseline['expression_score'],
            direction='decrease'
        )
        if facial.pupil_response_latency:
            deviations['pupil_latency'] = self._deviation(
                facial.pupil_response_latency, 
                self.baseline['pupil_latency'],
                direction='increase'
            )
        
        # 行为特征偏离
        deviations['steering'] = self._deviation(
            behavioral.steering_reversal_rate,
            self.baseline['steering_rate'],
            direction='decrease'  # 修正减少 = 损伤
        )
        deviations['lane_dev'] = self._deviation(
            behavioral.lane_deviation_std,
            self.baseline['lane_dev_std'],
            direction='increase'
        )
        deviations['speed_var'] = self._deviation(
            behavioral.speed_variation,
            self.baseline['speed_var'],
            direction='increase'
        )
        deviations['brake_rt'] = self._deviation(
            behavioral.brake_reaction_time,
            self.baseline['brake_rt'],
            direction='increase'
        )
        
        # 加权融合
        facial_score = sum(
            self.facial_weights[k] * deviations.get(k, 0)
            for k in self.facial_weights
        )
        behavioral_score = sum(
            self.behavioral_weights[k] * deviations.get(k, 0)
            for k in self.behavioral_weights
        )
        
        # 总损伤分数 (0-1)
        impairment = 0.5 * facial_score + 0.5 * behavioral_score
        impairment = np.clip(impairment, 0, 1)
        
        # BAC 估计
        bac_estimate = impairment * 0.12  # 线性映射到 BAC
        
        # 主要贡献指标
        sorted_dev = sorted(deviations.items(), key=lambda x: -x[1])
        primary = [k for k, v in sorted_dev[:3] if v > 0.3]
        
        # 建议
        if bac_estimate < self.BAC_WARNING:
            recommendation = 'safe'
        elif bac_estimate < self.BAC_UNSAFE:
            recommendation = 'warning'
        else:
            recommendation = 'unsafe'
        
        return SobrietyAssessment(
            impairment_score=impairment,
            bac_estimate=bac_estimate,
            confidence=min(impairment * 2, 1.0),
            primary_indicators=primary,
            recommendation=recommendation
        )
    
    def _deviation(self, value: float, baseline: float, 
                    direction: str = 'increase') -> float:
        """计算相对基线的偏离程度"""
        if direction == 'increase':
            dev = (value - baseline) / (baseline + 1e-6)
        else:
            dev = (baseline - value) / (baseline + 1e-6)
        return np.clip(dev, 0, 1)


# 测试
if __name__ == "__main__":
    detector = VisionBasedSobrietyDetector()
    
    # 场景1: 清醒驾驶员
    facial_sober = FacialCues(
        perclos=0.08, head_sway_amplitude=2.0, head_sway_frequency=0.3,
        facial_expression_score=0.6, pupil_response_latency=250
    )
    behavior_sober = BehavioralCues(
        steering_reversal_rate=30, lane_deviation_std=0.15,
        speed_variation=2.0, brake_reaction_time=0.8
    )
    r1 = detector.assess(facial_sober, behavior_sober)
    print(f"清醒: BAC估计={r1.bac_estimate:.3f}%, 建议={r1.recommendation}")
    
    # 场景2: BAC=0.08% 醉酒
    facial_drunk = FacialCues(
        perclos=0.16,       # +100%
        head_sway_amplitude=4.5,  # +125%
        head_sway_frequency=0.5,
        facial_expression_score=0.3,  # -50%
        pupil_response_latency=350   # +100ms
    )
    behavior_drunk = BehavioralCues(
        steering_reversal_rate=15,  # -50%
        lane_deviation_std=0.28,   # +87%
        speed_variation=4.5,       # +125%
        brake_reaction_time=1.2    # +50%
    )
    r2 = detector.assess(facial_drunk, behavior_drunk)
    print(f"醉酒: BAC估计={r2.bac_estimate:.3f}%, 建议={r2.recommendation}")
    print(f"  损伤分数: {r2.impairment_score:.2f}")
    print(f"  主要指标: {r2.primary_indicators}")

# 酒驾检测数据采集框架
class AlcoholImpairmentDataCollector:
    """酒精损伤数据采集器"""
    
    SCENARIOS = [
        # (场景, BAC目标, 时长min, 被试数)
        ('baseline_sober', 0.00, 30, 20),    # 清醒基线
        ('mild_impairment', 0.05, 30, 15),   # 轻度损伤
        ('moderate_impairment', 0.08, 30, 15), # 中度损伤
        ('severe_impairment', 0.12, 30, 10),  # 重度损伤
        ('fatigue_no_alcohol', 0.00, 60, 10), # 纯疲劳（对照组）
        ('mild_fatigue_mild_alcohol', 0.05, 45, 10), # 混合
    ]
    
    COLLECT_SIGNALS = [
        'camera_facial',      # 面部视频
        'camera_eye',         # 眼动
        'can_steering',       # 方向盘
        'can_speed',          # 速度
        'can_lane_position',  # 车道位置
        'can_brake',          # 刹车
        'simulator_hazard',   # 危险事件标记
    ]
    
    def get_dataset_stats(self):
        total_clips = sum(n * (t * 60 / 10) for _, _, t, n in self.SCENARIOS)
        return {
            'total_clips': int(total_clips),
            'total_subjects': sum(n for *_, n in self.SCENARIOS),
            'scenarios': len(self.SCENARIOS),
            'signals': len(self.COLLECT_SIGNALS)
        }