import numpy as np
import torch

def calculate_ela(landmarks_3d: np.ndarray) -> float:
    """
    计算 Eyelid Angle (ELA)
    
    论文 Section III-A
    基于 MediaPipe Face Mesh 3D landmarks
    
    Args:
        landmarks_3d: (468, 3) MediaPipe 3D 面部 landmarks
    
    Returns:
        ela: float, 眼睑角度（度），0=完全闭眼，~180=完全睁眼
    """
    # MediaPipe 眼睛 landmarks 索引
    # 左眼
    L_UPPER = [159, 158, 157, 173]  # 上眼睑
    L_LOWER = [145, 144, 153, 154]  # 下眼睑
    L_CORNER_L = 33   # 内眼角
    L_CORNER_R = 133  # 外眼角
    
    # 取上眼睑和下眼睑中心点
    upper_lid = np.mean(landmarks_3d[L_UPPER], axis=0)  # (3,)
    lower_lid = np.mean(landmarks_3d[L_LOWER], axis=0)  # (3,)
    corner_l = landmarks_3d[L_CORNER_L]  # (3,)
    corner_r = landmarks_3d[L_CORNER_R]  # (3,)
    
    # 计算眼睑方向向量
    lid_vector = lower_lid - upper_lid  # 上→下
    
    # 计算水平方向（眼角连线）
    horizontal = corner_r - corner_l
    
    # ELA = 眼睑向量与水平方向的夹角
    cos_angle = np.dot(lid_vector, horizontal) / (
        np.linalg.norm(lid_vector) * np.linalg.norm(horizontal) + 1e-8
    )
    ela = np.degrees(np.arccos(np.clip(cos_angle, -1, 1)))
    
    return ela


def calculate_ear(landmarks_2d: np.ndarray) -> float:
    """
    传统 EAR（对比基准）
    基于 2D landmarks 的眼睛纵横比
    
    Args:
        landmarks_2d: (6, 2) 眼睛 6 点 landmarks
    Returns:
        ear: float, EAR 值
    """
    # 6个点: 外角、上1、上2、内角、下1、下2
    p1, p2, p3, p4, p5, p6 = landmarks_2d
    
    # 垂直距离
    v1 = np.linalg.norm(p2 - p6)
    v2 = np.linalg.norm(p3 - p5)
    
    # 水平距离
    h = np.linalg.norm(p1 - p4)
    
    ear = (v1 + v2) / (2.0 * h + 1e-8)
    return ear


# ==================== 视角鲁棒性测试 ====================

def test_viewpoint_robustness():
    """对比 ELA 和 EAR 在不同头部角度下的稳定性"""
    print("=" * 60)
    print("ELA vs EAR 视角鲁棒性测试")
    print("=" * 60)
    
    # 模拟正常睁眼的 3D landmarks
    np.random.seed(42)
    base_landmarks = np.random.randn(468, 3) * 0.01
    # 设置眼睛区域
    base_landmarks[159] = [0.1, 0.05, 0.02]   # 上眼睑
    base_landmarks[145] = [0.1, -0.05, 0.02]  # 下眼睑
    base_landmarks[33] = [0.05, 0.0, 0.0]     # 内角
    base_landmarks[133] = [0.15, 0.0, 0.0]    # 外角
    
    rotations = {
        '正面 (0°)': np.array([0, 0, 0]),
        '偏航 15°': np.array([0, 15, 0]),
        '偏航 30°': np.array([0, 30, 0]),
        '俯仰 15°': np.array([15, 0, 0]),
        '俯仰 -15°': np.array([-15, 0, 0]),
        '滚转 10°': np.array([0, 0, 10]),
    }
    
    print(f"{'角度':15s} {'ELA (°)':>10s} {'EAR':>10s} {'ELA变化':>10s} {'EAR变化':>10s}")
    print("-" * 60)
    
    base_ela = calculate_ela(base_landmarks)
    base_2d = base_landmarks[:, :2]
    base_ear = calculate_ear(base_2d[[33, 159, 158, 133, 145, 153]])
    
    for name, angles in rotations.items():
        # 旋转 landmarks
        rx, ry, rz = np.radians(angles)
        Rx = np.array([[1, 0, 0], [0, np.cos(rx), -np.sin(rx)], [0, np.sin(rx), np.cos(rx)]])
        Ry = np.array([[np.cos(ry), 0, np.sin(ry)], [0, 1, 0], [-np.sin(ry), 0, np.cos(ry)]])
        Rz = np.array([[np.cos(rz), -np.sin(rz), 0], [np.sin(rz), np.cos(rz), 0], [0, 0, 1]])
        R = Rx @ Ry @ Rz
        
        rotated = (base_landmarks @ R.T)
        
        ela = calculate_ela(rotated)
        ear = calculate_ear(rotated[:, :2][[33, 159, 158, 133, 145, 153]])
        
        ela_change = abs(ela - base_ela)
        ear_change = abs(ear - base_ear) / base_ear * 100
        
        print(f"{name:15s} {ela:10.2f} {ear:10.4f} {ela_change:10.2f}° {ear_change:9.1f}%")
    
    print("\n结论: ELA 在头部旋转下变化远小于 EAR")


# ==================== 眨眼检测 ====================

class BlinkDetector:
    """
    基于 ELA 的眨眼检测器
    
    论文 Section III-B
    提取眨眼时间特征用于嗜睡判断
    """
    
    def __init__(self, ela_threshold: float = 120.0,
                 min_duration: float = 0.03,
                 max_duration: float = 2.0,
                 fps: int = 30):
        """
        Args:
            ela_threshold: 闭眼 ELA 阈值（度）
            min_duration: 最小眨眼时长（秒）
            max_duration: 最大眨眼时长（秒），超过视为微睡眠
            fps: 帧率
        """
        self.threshold = ela_threshold
        self.min_frames = int(min_duration * fps)
        self.max_frames = int(max_duration * fps)
        self.fps = fps
        
    def detect_blinks(self, ela_sequence: np.ndarray) -> list:
        """
        检测眨眼事件
        
        Args:
            ela_sequence: (N,) ELA 时间序列
        Returns:
            blinks: list of {start, end, closing_dur, closed_dur, reopening_dur, total_dur}
        """
        closed = ela_sequence < self.threshold
        
        blinks = []
        i = 0
        while i < len(closed):
            if closed[i]:
                # 找到闭眼开始
                start = i
                while i < len(closed) and closed[i]:
                    i += 1
                end = i
                
                duration = (end - start) / self.fps
                
                if duration >= self.min_frames / self.fps:
                    # 找到 reopen 结束（ELA 回到阈值以上）
                    reopen_start = end
                    while i < len(closed) and not closed[i]:
                        i += 1
                    reopen_end = i if i < len(closed) else len(closed)
                    
                    # 分析前闭眼阶段（closing）
                    pre_start = max(0, start - (reopen_end - reopen_start))
                    
                    blinks.append({
                        'start_frame': start,
                        'end_frame': end,
                        'reopen_end': reopen_end,
                        'closing_duration': (end - start) / (2 * self.fps),  # 估算
                        'closed_duration': duration,
                        'reopening_duration': (reopen_end - reopen_start) / self.fps,
                        'total_duration': (reopen_end - start) / self.fps,
                        'is_microsleep': duration > self.max_frames / self.fps
                    })
            else:
                i += 1
        
        return blinks
    
    def extract_drowsiness_features(self, blinks: list, 
                                     window_sec: int = 60) -> dict:
        """
        提取嗜睡特征
        
        Args:
            blinks: 眨眼事件列表
            window_sec: 分析窗口（秒）
        Returns:
            features: 嗜睡特征字典
        """
        if not blinks:
            return {'drowsiness_score': 0}
        
        closed_durations = [b['closed_duration'] for b in blinks]
        reopening_durations = [b['reopening_duration'] for b in blinks]
        total_durations = [b['total_duration'] for b in blinks]
        microsleeps = sum(1 for b in blinks if b['is_microsleep'])
        long_blinks = sum(1 for d in closed_durations if d > 0.15)  # >150ms
        
        # 嗜睡评分（0-100）
        avg_closed = np.mean(closed_durations)
        avg_reopen = np.mean(reopening_durations)
        long_ratio = long_blinks / len(blinks)
        
        score = (avg_closed / 0.3 * 30 +  # 闭眼时长权重
                 avg_reopen / 0.5 * 30 +   # reopen 时长权重
                 long_ratio * 100 * 25 +   # 长眨眼比例权重
                 microsleeps * 15)         # 微睡眠权重
        
        return {
            'blink_count': len(blinks),
            'avg_closed_duration': avg_closed,
            'avg_reopening_duration': avg_reopen,
            'long_blink_ratio': long_ratio,
            'microsleep_count': microsleeps,
            'drowsiness_score': min(score, 100),
            'drowsiness_level': self._classify_level(score)
        }
    
    def _classify_level(self, score: float) -> str:
        if score < 25: return 'ALERT'
        elif score < 50: return 'SLIGHTLY_DROWSY'
        elif score < 75: return 'DROWSY'
        else: return 'VERY_DROWSY'


# ==================== 合成数据生成 ====================

class SyntheticBlinkGenerator:
    """
    基于 ELA 的合成眨眼数据生成
    
    论文 Section IV
    用 ELA 信号驱动 Blender 3D 虚拟人
    """
    
    @staticmethod
    def generate_blink_ela(duration_sec: float = 1.0,
                           fps: int = 30,
                           blink_rate: float = 15.0,
                           drowsiness: float = 0.0) -> np.ndarray:
        """
        生成模拟 ELA 序列
        
        Args:
            duration_sec: 总时长
            fps: 帧率
            blink_rate: 每分钟眨眼次数
            drowsiness: 嗜睡程度 0-1
        Returns:
            ela_sequence: (N,) 模拟 ELA 值
        """
        n_frames = int(duration_sec * fps)
        n_blinks = int(duration_sec * blink_rate / 60)
        
        # 基线 ELA（睁眼）
        ela = np.full(n_frames, 160.0)
        
        # 添加噪声
        ela += np.random.normal(0, 2, n_frames)
        
        for _ in range(n_blinks):
            # 随机眨眼时间
            start = np.random.randint(0, n_frames - 60)
            
            # 嗜睡影响：闭眼更久、reopen 更慢
            closed_dur = int((0.05 + drowsiness * 0.2) * fps)
            reopen_dur = int((0.1 + drowsiness * 0.4) * fps)
            closing_dur = int(0.03 * fps)
            
            # 闭合阶段
            for j in range(closing_dur):
                if start + j < n_frames:
                    progress = j / closing_dur
                    ela[start + j] = 160 * (1 - progress) + 30 * progress
            
            # 闭眼阶段
            for j in range(closed_dur):
                if start + closing_dur + j < n_frames:
                    ela[start + closing_dur + j] = 30
            
            # reopen 阶段
            for j in range(reopen_dur):
                if start + closing_dur + closed_dur + j < n_frames:
                    progress = j / reopen_dur
                    ela[start + closing_dur + closed_dur + j] = 30 + 130 * progress
        
        return ela


# ==================== 完整测试 ====================

if __name__ == "__main__":
    # 1. 视角鲁棒性测试
    test_viewpoint_robustness()
    
    # 2. 眨眼检测测试
    print("\n" + "=" * 60)
    print("眨眼检测与嗜睡评分测试")
    print("=" * 60)
    
    # 生成模拟数据
    alert_ela = SyntheticBlinkGenerator.generate_blink(60, drowsiness=0.0)
    drowsy_ela = SyntheticBlinkGenerator.generate_blink(60, drowsiness=0.8)
    
    detector = BlinkDetector(fps=30)
    
    alert_blinks = detector.detect_blinks(alert_ela)
    drowsy_blinks = detector.detect_blinks(drowsy_ela)
    
    alert_features = detector.extract_drowsiness_features(alert_blinks)
    drowsy_features = detector.extract_drowsiness_features(drowsy_blinks)
    
    print("\n清醒状态:")
    for k, v in alert_features.items():
        print(f"  {k}: {v}")
    
    print("\n嗜睡状态:")
    for k, v in drowsy_features.items():
        print(f"  {k}: {v}")
    
    # 3. 对比 EAR vs ELA
    print("\n" + "=" * 60)
    print("EAR vs ELA 稳定性对比")
    print("=" * 60)
    
    # 模拟 1000 帧
    np.random.seed(0)
    ela_values = SyntheticBlinkGenerator.generate_blink(33.3, drowsiness=0.3)
    
    print(f"ELA 统计: 均值={np.mean(ela_values):.2f}°, 标准差={np.std(ela_values):.2f}°")
    print(f"ELA 范围: [{np.min(ela_values):.1f}°, {np.max(ela_values):.1f}°]")
    print(f"眨眼事件数: {len(alert_blinks)} (清醒) vs {len(drowsy_blinks)} (嗜睡)")

class DMSDrowsinessUpgrade:
    """从 EAR 升级到 ELA 的实施方案"""
    
    def __init__(self):
        self.upgrade_path = {
            'sensor': '无需更换（同一摄像头）',
            'landmark_model': 'MediaPipe Face Mesh（已部署）',
            'metric_change': 'EAR → ELA',
            'model_retrain': '需重训嗜睡分类器',
            'expected_gain': '+9.4% 准确率，+19×视角鲁棒性'
        }
    
    def implementation_steps(self):
        return [
            '1. 替换 metric 计算函数（EAR→ELA）',
            '2. 重训眨眼检测阈值（ELA 角度域）',
            '3. 重训嗜睡分类器（加入新特征）',
            '4. 用 Blender 合成数据增强训练集',
            '5. 验证：真实数据 + 合成数据联合训练'
        ]