import numpy as np
from datetime import datetime, timedelta

class AviationFatigueModel:
    """
    航空疲劳预测模型（简化版）
    
    基于三个维度：
    1. 生物节律（昼夜节律）
    2. 持续清醒时间
    3. 任务负荷
    
    这三个维度同样适用于汽车驾驶员
    """
    
    def __init__(self):
        # 昼夜节律参数（人体核心体温周期）
        self.circadian_phase = 0  # 0-24小时
        
    def circadian_alertness(self, hour: float) -> float:
        """
        昼夜节律驱动的警觉度
        
        基于人体核心体温周期：
        - 体温最高点（~18:00）→ 最高警觉
        - 体温最低点（~05:00）→ 最低警觉
        
        Returns:
            alertness: 0-1, 1=最高警觉
        """
        # 余弦模型，峰值在下午6点
        phase = (hour - 18) * 2 * np.pi / 24
        alertness = 0.5 + 0.3 * np.cos(phase)
        return max(0.1, min(1.0, alertness))
    
    def sleep_pressure(self, hours_awake: float) -> float:
        """
        持续清醒的睡眠压力
        
        航空医学：清醒>16小时≈BAC 0.05%
        清醒>20小时≈BAC 0.10%
        
        Returns:
            sleep_pressure: 0-1, 1=最大压力
        """
        # 指数增长模型
        pressure = 1 - np.exp(-hours_awake / 12)
        return min(1.0, pressure)
    
    def workload_factor(self, task_complexity: float,
                        duration_hours: float) -> float:
        """
        任务负荷因子
        
        航空：巡航(低负荷) vs 起降(高负荷)
        汽车：高速巡航 vs 城市驾驶
        
        Args:
            task_complexity: 0-1 (0=巡航, 1=复杂)
            duration_hours: 持续时间
        """
        # 持续高负荷→更快疲劳
        fatigue_rate = 0.05 + 0.15 * task_complexity
        workload_fatigue = fatigue_rate * duration_hours
        return min(1.0, workload_fatigue)
    
    def predict_fatigue(self, hour: float, hours_awake: float,
                        task_complexity: float,
                        duration: float) -> dict:
        """
        综合疲劳预测
        
        Returns:
            fatigue_level: 0-1
            risk_level: low/moderate/high/severe
            primary_factor: 主要疲劳因素
        """
        circadian = self.circadian_alertness(hour)
        sleep_p = self.sleep_pressure(hours_awake)
        workload = self.workload_factor(task_complexity, duration)
        
        # 综合疲劳 = (1-警觉) × 0.3 + 睡眠压力 × 0.4 + 工作负荷 × 0.3
        fatigue = (
            (1 - circadian) * 0.3 +
            sleep_p * 0.4 +
            workload * 0.3
        )
        
        # 分级
        if fatigue < 0.3:
            level = 'low'
        elif fatigue < 0.5:
            level = 'moderate'
        elif fatigue < 0.7:
            level = 'high'
        else:
            level = 'severe'
        
        # 主要因素
        factors = {
            'circadian': (1 - circadian) * 0.3,
            'sleep_pressure': sleep_p * 0.4,
            'workload': workload * 0.3
        }
        primary = max(factors, key=factors.get)
        
        return {
            'fatigue_level': fatigue,
            'risk_level': level,
            'primary_factor': primary,
            'circadian_alertness': circadian,
            'sleep_pressure': sleep_p,
            'workload_fatigue': workload
        }


# 测试：航空 vs 汽车场景对比
if __name__ == "__main__":
    model = AviationFatigueModel()
    
    scenarios = [
        ("汽车-日间短途", 14, 4, 0.3, 1),
        ("汽车-夜间长途", 2, 8, 0.5, 4),
        ("汽车-凌晨通勤", 6, 16, 0.4, 0.5),
        ("航空-跨时区航班", 3, 14, 0.6, 8),
        ("铁路-夜班固定线", 1, 10, 0.2, 6),
    ]
    
    print(f"{'场景':20s} | {'时间':>4} | {'清醒':>4} | {'负荷':>4} | {'时长':>4} | "
          f"{'疲劳':>6} | {'风险':>10} | {'主因'}")
    print("-" * 90)
    
    for name, hour, awake, complexity, duration in scenarios:
        result = model.predict_fatigue(hour, awake, complexity, duration)
        print(f"{name:20s} | {hour:4.0f}h | {awake:4.0f}h | {complexity:4.1f} | "
              f"{duration:4.1f}h | {result['fatigue_level']:.3f} | "
              f"{result['risk_level']:>10} | {result['primary_factor']}")

# 将昼夜节律集成到现有DMS中
class CircadianEnhancedDMS:
    """
    航空疲劳预测模型增强版DMS
    
    在传统DMS基础上增加：
    1. 出发时间感知
    2. 昼夜节律疲劳基线
    3. 持续清醒时间估计
    4. 任务负荷自适应阈值
    """
    
    def __init__(self):
        self.fatigue_model = AviationFatigueModel()
        self.perceived_fatigue = 0  # DMS检测值
        self.predicted_fatigue = 0  # 模型预测值
        
    def update_context(self, current_time, hours_awake,
                       driving_duration, road_type):
        """更新驾驶上下文"""
        self.current_hour = current_time
        self.hours_awake = hours_awake
        self.driving_duration = driving_duration
        
        # 任务复杂度
        complexity_map = {
            'highway_cruise': 0.2,
            'city_traffic': 0.6,
            'night_rain': 0.8
        }
        self.complexity = complexity_map.get(road_type, 0.5)
        
        # 预测疲劳
        result = self.fatigue_model.predict_fatigue(
            current_time, hours_awake,
            self.complexity, driving_duration
        )
        self.predicted_fatigue = result['fatigue_level']
        self.fatigue_primary = result['primary_factor']
    
    def evaluate(self, dms_detection_result):
        """
        融合DMS检测结果和预测模型
        """
        # DMS实测值（PERCLOS等）
        self.perceived_fatigue = dms_detection_result.get('perclos_normalized', 0)
        
        # 融合：实测权重0.7，预测权重0.3
        fused = 0.7 * self.perceived_fatigue + 0.3 * self.predicted_fatigue
        
        # 自适应阈值
        base_threshold = 0.5
        # 预测高疲劳→降低阈值（更敏感）
        if self.predicted_fatigue > 0.6:
            threshold = base_threshold - 0.1
        else:
            threshold = base_threshold
        
        return {
            'fatigue_level': fused,
            'threshold': threshold,
            'alert': fused > threshold,
            'primary_cause': self.fatigue_primary,
            'predicted': self.predicted_fatigue,
            'perceived': self.perceived_fatigue
        }


# 测试
if __name__ == "__main__":
    dms = CircadianEnhancedDMS()
    
    # 凌晨2点，清醒18小时，驾驶4小时，夜间高速
    dms.update_context(
        current_time=2,
        hours_awake=18,
        driving_duration=4,
        road_type='highway_cruise'
    )
    
    # DMS实测：PERCLOS不高（DMS说还好）
    result = dms.evaluate({'perclos_normalized': 0.3})
    
    print("凌晨2点+清醒18小时场景:")
    print(f"  DMS实测: {result['perceived']:.3f}")
    print(f"  模型预测: {result['predicted']:.3f}")
    print(f"  融合结果: {result['fatigue_level']:.3f}")
    print(f"  自适应阈值: {result['threshold']:.3f}")
    print(f"  触发警报: {result['alert']}")
    print(f"  主要原因: {result['primary_cause']}")