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

@dataclass
class NASA SensorSpec:
    """NASA 乘员监测传感器规格"""
    name: str
    modality: str        # 生理/行为
    placement: str       # 佩戴位置
    non_obtrusive: bool  # 非侵入性
    power_mw: float      # 功耗 (mW)
    sampling_rate: int   # 采样率 (Hz)

# NASA 推荐的非侵入式传感器套件
NASA_SENSOR_SUITE = [
    NASASensorSpec("ECG/HRV", "生理", "胸带/背心", True, 5, 250),
    NASASensorSpec("眼动追踪", "行为", "头戴/固定", True, 100, 60),
    NASASensorSpec("皮电响应", "生理", "腕带", True, 2, 32),
    NASASensorSpec("血氧SpO2", "生理", "耳夹/腕带", True, 10, 1),
    NASASensorSpec("体温", "生理", "贴片", True, 1, 1),
    NASASensorSpec("EEG(简化)", "生理", "头带", True, 50, 250),
]

class CrewStateClassifier:
    """
    NASA 乘员状态分类器
    
    基于心理生理监测的群体分类器 + 个性化适配
    
    流程：
    1. 群体基线建立（任务前）
    2. 实时状态分类（任务中）
    3. 个性化模型更新（长期）
    4. 对策建议（闭环干预）
    """
    
    def __init__(self, sensor_specs: List[NASA SensorSpec] = None):
        self.sensors = sensor_specs or NASA_SENSOR_SUITE
        self.baseline = None
        self.personalized = False
        self.alert_history = []
        
    def establish_baseline(
        self,
        rest_data: dict,
        duration_min: int = 10
    ) -> dict:
        """
        建立群体基线
        
        Args:
            rest_data: 静息状态传感器数据
            duration_min: 基线采集时长
            
        Returns:
            baseline: 基线指标
        """
        self.baseline = {
            'hrv_rmssd': np.mean(rest_data.get('hrv', [40])),
            'blink_rate': np.mean(rest_data.get('blink_rate', [18])),
            'gaze_entropy': np.mean(rest_data.get('gaze_entropy', [2.5])),
            'eda_baseline': np.mean(rest_data.get('eda', [3.0])),
            'spo2': np.mean(rest_data.get('spo2', [97])),
            'temp': np.mean(rest_data.get('temp', [36.8])),
        }
        return self.baseline
    
    def classify_state(self, current_data: dict) -> dict:
        """
        实时乘员状态分类
        
        Args:
            current_data: 当前传感器数据
            
        Returns:
            classification: 状态分类结果
        """
        if self.baseline is None:
            return {'error': '基线未建立'}
        
        # 计算偏差
        deviations = {}
        for key in self.baseline:
            if key in current_data:
                base = self.baseline[key]
                curr = np.mean(current_data[key])
                deviations[key] = (curr - base) / (base + 1e-8)
        
        # 状态判定规则
        states = []
        
        # 疲劳判定
        hrv_dev = deviations.get('hrv_rmssd', 0)
        blink_dev = deviations.get('blink_rate', 0)
        if hrv_dev < -0.2 and blink_dev > 0.3:
            states.append({
                'state': 'fatigue',
                'confidence': min(abs(hrv_dev) + abs(blink_dev), 1.0),
                'severity': 'moderate' if abs(hrv_dev) > 0.3 else 'mild'
            })
        
        # 压力判定
        eda_dev = deviations.get('eda_baseline', 0)
        if eda_dev > 0.3:
            states.append({
                'state': 'stress',
                'confidence': min(abs(eda_dev), 1.0),
                'severity': 'high' if eda_dev > 0.5 else 'moderate'
            })
        
        # 注意力失败
        gaze_dev = deviations.get('gaze_entropy', 0)
        if gaze_dev < -0.15:
            states.append({
                'state': 'channelized_attention',
                'confidence': min(abs(gaze_dev), 1.0),
                'severity': 'moderate'
            })
        elif gaze_dev > 0.3:
            states.append({
                'state': 'diverted_attention',
                'confidence': min(abs(gaze_dev), 1.0),
                'severity': 'moderate'
            })
        
        # 综合风险评估
        total_risk = sum(s['confidence'] for s in states)
        
        return {
            'states': states,
            'total_risk': min(total_risk, 1.0),
            'recommendation': self._get_countermeasure(states, total_risk),
            'deviations': deviations
        }
    
    def _get_countermeasure(self, states: list, risk: float) -> str:
        """获取对策建议"""
        if risk < 0.2:
            return 'nominal'
        elif risk < 0.5:
            return 'increase_automation'
        elif risk < 0.8:
            return 'task_redistribution'
        else:
            return 'emergency_protocol'


if __name__ == "__main__":
    classifier = CrewStateClassifier()
    
    print("=== NASA 乘员监测传感器套件 ===")
    for s in classifier.sensors:
        print(f"{s.name:15s} | {s.modality:4s} | {s.placement:12s} | "
              f"{'非侵入' if s.non_obtrusive else '侵入':4s} | "
              f"{s.power_mw:5.1f}mW | {s.sampling_rate:4d}Hz")
    
    # 建立基线
    rest_data = {
        'hrv': [42, 40, 38, 41, 43],
        'blink_rate': [17, 19, 18, 16, 18],
        'gaze_entropy': [2.6, 2.4, 2.5, 2.7, 2.5],
        'eda': [2.8, 3.0, 3.1, 2.9, 3.0],
        'spo2': [97, 98, 97, 98, 98],
        'temp': [36.7, 36.8, 36.9, 36.8, 36.8],
    }
    baseline = classifier.establish_baseline(rest_data)
    print(f"\n基线: {baseline}")
    
    # 疲劳状态
    fatigued_data = {
        'hrv': [30, 28, 32, 29],
        'blink_rate': [25, 28, 24, 26],
        'gaze_entropy': [2.0, 1.8, 2.1, 1.9],
        'eda': [3.5, 3.8, 3.6, 3.7],
    }
    result = classifier.classify_state(fatigued_data)
    print(f"\n=== 疲劳状态分析 ===")
    for s in result.get('states', []):
        print(f"状态: {s['state']}, 置信度: {s['confidence']:.2f}, 严重度: {s['severity']}")
    print(f"总风险: {result['total_risk']:.2f}")
    print(f"对策: {result['recommendation']}")

class IMSCrewStateMonitor:
    """
    IMS 乘员状态监测器（受NASA启发）
    
    从二分类(疲劳/正常)升级到多状态分类
    """
    
    STATES = [
        'normal',           # 正常驾驶
        'mild_fatigue',     # 轻度疲劳
        'moderate_fatigue', # 中度疲劳
        'cognitive_distraction',  # 认知分心
        'visual_distraction',     # 视觉分心
        'stress',           # 压力
        'impairment',       # 损伤(酒精/药物)
    ]
    
    def __init__(self, config: dict = None):
        self.thresholds = config or {
            'normal': 0.0,
            'alert': 0.3,      # 一级告警
            'warning': 0.6,    # 二级警告
            'critical': 0.85,  # 三级紧急
        }
        
    def assess(self, features: dict) -> dict:
        """多状态评估"""
        scores = {}
        
        # 疲劳评分
        perclos = features.get('perclos', 0)
        blink_rate = features.get('blink_rate', 18)
        scores['fatigue'] = (
            0.5 * min(perclos / 0.3, 1.0) +
            0.3 * min(max(0, (blink_rate - 20)) / 15, 1.0) +
            0.2 * features.get('yawn_score', 0)
        )
        
        # 分心评分
        gaze_deviation = features.get('gaze_deviation', 0)
        gaze_entropy = features.get('gaze_entropy', 0)
        scores['distraction'] = (
            0.5 * min(gaze_deviation / 0.3, 1.0) +
            0.5 * min(max(0, gaze_entropy - 2.5) / 1.0, 1.0)
        )
        
        # 压力评分（如果EDA可用）
        eda_change = features.get('eda_change', 0)
        hr_change = features.get('hr_change', 0)
        scores['stress'] = (
            0.5 * min(eda_change / 0.3, 1.0) +
            0.5 * min(hr_change / 20, 1.0)
        )
        
        # 确定主要状态
        primary = max(scores, key=scores.get)
        primary_score = scores[primary]
        
        # 告警等级
        if primary_score >= self.thresholds['critical']:
            level = 3
        elif primary_score >= self.thresholds['warning']:
            level = 2
        elif primary_score >= self.thresholds['alert']:
            level = 1
        else:
            level = 0
        
        return {
            'primary_state': primary,
            'scores': scores,
            'alert_level': level,
            'recommendation': self._recommend(level, primary)
        }
    
    def _recommend(self, level: int, state: str) -> str:
        if level == 0:
            return 'nominal'
        elif level == 1:
            return f'visual_audio_alert ({state})'
        elif level == 2:
            return f'haptic_alert + adas_assist ({state})'
        else:
            return f'safe_stop_protocol ({state})'