"""
Eyelid Angle (ELA) 指标实现

论文核心：从3D面部关键点计算上下眼睑之间的角度
相比EAR的2D距离比，ELA在头部旋转时保持稳定
"""

import numpy as np
from typing import Tuple, Optional

class EyelidAngle:
    """
    眼睑角度 (ELA) 计算
    
    基于 MediaPipe Face Mesh 3D 关键点：
    - 上眼睑关键点（6个）
    - 下眼睑关键点（6个）
    
    ELA = 上下眼睑切线之间的角度
    """
    
    # MediaPipe Face Mesh 眼部关键点索引
    LEFT_EYE_UPPER = [33, 246, 161, 160, 159, 158]  # 左眼上睑
    LEFT_EYE_LOWER = [145, 144, 153, 154, 155, 133]  # 左眼下睑
    RIGHT_EYE_UPPER = [362, 398, 384, 385, 386, 387]  # 右眼上睑
    RIGHT_EYE_LOWER = [263, 373, 374, 380, 381, 382]  # 右眼下睑
    
    def __init__(self):
        pass
    
    def compute_ela(self, landmarks_3d: np.ndarray, eye: str = 'left') -> float:
        """
        计算眼睑角度
        
        Args:
            landmarks_3d: (468, 3) MediaPipe 3D面部关键点
            eye: 'left' 或 'right'
            
        Returns:
            ela_angle: 眼睑角度（度），0°=完全闭合，>30°=完全睁开
        """
        if eye == 'left':
            upper_idx = self.LEFT_EYE_UPPER
            lower_idx = self.LEFT_EYE_LOWER
        else:
            upper_idx = self.RIGHT_EYE_UPPER
            lower_idx = self.RIGHT_EYE_LOWER
        
        # 提取上下眼睑关键点
        upper_pts = landmarks_3d[upper_idx]  # (6, 3)
        lower_pts = landmarks_3d[lower_idx]  # (6, 3)
        
        # 拟合上眼睑切线（最小二乘）
        upper_dir = self._fit_line_direction(upper_pts)
        
        # 拟合下眼睑切线
        lower_dir = self._fit_line_direction(lower_pts)
        
        # 计算两切线之间的角度
        angle = self _angle_between_vectors(upper_dir, lower_dir)
        
        return np.degrees(angle)
    
    def _fit_line_direction(self, points: np.ndarray) -> np.ndarray:
        """最小二乘拟合方向向量"""
        center = points.mean(axis=0)
        centered = points - center
        # SVD 分解取主方向
        _, _, vh = np.linalg.svd(centered)
        direction = vh[0]  # 最大奇异值方向
        return direction
    
    def _angle_between_vectors(self, v1: np.ndarray, v2: np.ndarray) -> float:
        """两3D向量之间的角度（弧度）"""
        cos_angle = np.dot(v1, v2) / (
            np.linalg.norm(v1) * np.linalg.norm(v2) + 1e-8
        )
        cos_angle = np.clip(cos_angle, -1, 1)
        return np.arccos(cos_angle)
    
    def compute_ear_baseline(self, landmarks_2d: np.ndarray, eye: str = 'left') -> float:
        """
        传统 EAR 基线（对比用）
        
        EAR = (|p2-p6| + |p3-p5|) / (2 * |p1-p4|)
        """
        if eye == 'left':
            p1, p2, p3, p4, p5, p6 = 33, 160, 158, 133, 153, 144
        else:
            p1, p2, p3, p4, p5, p6 = 362, 385, 387, 263, 373, 374
        
        pts = landmarks_2d[[p1, p2, p3, p4, p5, p6]]
        ear = (
            np.linalg.norm(pts[1] - pts[5]) +
            np.linalg.norm(pts[2] - pts[4])
        ) / (2 * np.linalg.norm(pts[0] - pts[3]))
        
        return ear
    
    def compute_blink_features(
        self,
        ela_sequence: np.ndarray,
        fps: int = 30,
        closing_threshold: float = 15.0,  # 度
        closed_threshold: float = 5.0     # 度
    ) -> dict:
        """
        从 ELA 时序信号提取眨眼特征
        
        Args:
            ela_sequence: ELA 角度序列（度），shape=(N,)
            fps: 帧率
            closing_threshold: 开始闭眼的 ELA 阈值
            closed_threshold: 完全闭眼的 ELA 阈值
            
        Returns:
            blink_features: 眨眼时间特征
        """
        features = []
        
        i = 0
        while i < len(ela_sequence) - 1:
            # 检测闭眼开始（ELA 下降到阈值以下）
            if ela_sequence[i] < closing_threshold and ela_sequence[i-1] >= closing_threshold:
                blink_start = i
                
                # 找到最低点
                j = i
                while j < len(ela_sequence) and ela_sequence[j] < closing_threshold:
                    j += 1
                
                if j < len(ela_sequence):
                    # 闭眼持续时间
                    closed_duration = (j - blink_start) / fps * 1000  # ms
                    
                    # 找到完全闭眼段
                    closed_start = blink_start
                    closed_end = j
                    while closed_start < j and ela_sequence[closed_start] > closed_threshold:
                        closed_start += 1
                    while closed_end > blink_start and ela_sequence[closed_end] > closed_threshold:
                        closed_end -= 1
                    
                    fully_closed = (closed_end - closed_start) / fps * 1000 if closed_end > closed_start else 0
                    
                    # 重新睁眼时间
                    reopen_start = j
                    while reopen_start < len(ela_sequence) and ela_sequence[reopen_start] < closing_threshold:
                        reopen_start += 1
                    reopen_duration = (reopen_start - j) / fps * 1000
                    
                    # 总眨眼时长
                    total_duration = (reopen_start - blink_start) / fps * 1000
                    
                    # 幅度
                    amplitude = closing_threshold - np.min(ela_sequence[blink_start:j+1])
                    
                    features.append({
                        'closing_duration': (j - blink_start) / fps * 1000,  # ms
                        'closed_duration': fully_closed,  # ms
                        'reopen_duration': reopen_duration,  # ms
                        'total_duration': total_duration,  # ms
                        'amplitude': amplitude,  # 度
                        'amp_vel_ratio': amplitude / ((j - blink_start) / fps + 1e-8),  # 度/秒
                    })
                    
                    i = reopen_start
                else:
                    i = j
            else:
                i += 1
        
        if not features:
            return {'n_blinks': 0, 'avg_duration': 0, 'long_blink_ratio': 0}
        
        durations = [f['total_duration'] for f in features]
        long_blinks = sum(1 for d in durations if d > 500)  # >500ms 为长眨眼
        
        return {
            'n_blinks': len(features),
            'avg_closing': np.mean([f['closing_duration'] for f in features]),
            'avg_closed': np.mean([f['closed_duration'] for f in features]),
            'avg_reopen': np.mean([f['reopen_duration'] for f in features]),
            'avg_total': np.mean(durations),
            'long_blink_ratio': long_blinks / len(features),
            'avg_amplitude': np.mean([f['amplitude'] for f in features]),
            'perclos_long': sum(1 for d in durations if d > 500) / len(ela_sequence) * fps * 100,
            'features': features
        }


# ============ 测试 ============
if __name__ == "__main__":
    np.random.seed(42)
    
    ela = EyelidAngle()
    
    # 模拟3D面部关键点 (468个点, 3D坐标)
    landmarks = np.random.randn(468, 3) * 100
    
    # 设置眼部关键点（模拟睁眼状态）
    for idx in ela.LEFT_EYE_UPPER:
        landmarks[idx] = [0, 0.5, -5]  # 上眼睑位置
    for idx in ela.LEFT_EYE_LOWER:
        landmarks[idx] = [0, -0.5, -5]  # 下眼睑位置
    
    # 计算ELA
    angle = ela.compute_ela(landmarks, eye='left')
    print(f"睁眼状态 ELA: {angle:.1f}°")
    
    # 模拟闭眼（上眼睑下移）
    for idx in ela.LEFT_EYE_UPPER:
        landmarks[idx] = [0, -0.3, -5]
    angle_closed = ela.compute_ela(landmarks, eye='left')
    print(f"闭眼状态 ELA: {angle_closed:.1f}°")
    
    # 生成眨眼序列
    n_frames = 150  # 5秒 @ 30fps
    ela_seq = np.full(n_frames, 30.0)  # 睁眼基线
    ela_seq[50:60] = np.linspace(30, 2, 10)  # 闭眼
    ela_seq[60:68] = 2.0  # 闭眼持续
    ela_seq[68:80] = np.linspace(2, 30, 12)  # 睁眼
    
    # 提取眨眼特征
    blink = ela.compute_blink_features(ela_seq, fps=30)
    print(f"\n眨眼特征:")
    print(f"  眨眼次数: {blink['n_blinks']}")
    print(f"  闭眼时长: {blink['avg_closed']:.0f}ms")
    print(f"  重新睁眼: {blink['avg_reopen']:.0f}ms")
    print(f"  总时长: {blink['avg_total']:.0f}ms")
    print(f"  长眨眼比: {blink['long_blink_ratio']:.1%}")

"""
Blender 3D 合成眨眼数据管道

论文 Section III-C: 使用ELA信号驱动Blender中的rigged avatar
生成可控噪声、视角、眨眼动态的真实合成数据
"""

class SyntheticBlinkGenerator:
    """
    Blender 驱动的合成眨眼数据生成
    
    优势：
    1. 无需危险的真实疲劳驾驶数据
    2. 可控制眨眼参数
    3. 多视角生成
    4. 噪声注入增强鲁棒性
    """
    
    def __init__(self):
        self.parameters = {
            'avatar_models': ['male_01', 'female_01', 'male_02', 'female_02'],
            'camera_positions': [
                ('front', (0, 0, 30)),
                ('upper_left', (-15, 10, 30)),
                ('upper_right', (15, 10, 30)),
                ('side_left', (-30, 0, 25)),
                ('side_right', (30, 0, 25)),
            ],
            'lighting': ['daylight', 'tunnel', 'night_ir', 'mixed'],
            'noise_levels': [0, 5, 10, 20],  # dB
        }
        
        # 疲劳等级对应眨眼参数
        self.fatigue_levels = {
            'alert': {
                'blink_rate': (10, 20),  # 次/分钟
                'closing_duration': (50, 100),  # ms
                'closed_duration': (50, 100),
                'reopen_duration': (100, 200),
                'amplitude': (25, 35),  # 度
            },
            'low_vigilant': {
                'blink_rate': (15, 25),
                'closing_duration': (80, 150),
                'closed_duration': (80, 150),
                'reopen_duration': (150, 300),
                'amplitude': (20, 30),
            },
            'drowsy': {
                'blink_rate': (20, 35),
                'closing_duration': (150, 300),
                'closed_duration': (200, 500),
                'reopen_duration': (300, 600),
                'amplitude': (15, 25),
            },
            'microsleep': {
                'blink_rate': (5, 15),  # 频率低但时长长
                'closing_duration': (300, 500),
                'closed_duration': (500, 2000),
                'reopen_duration': (500, 1000),
                'amplitude': (10, 20),
            }
        }
    
    def generate_config(self, level: str = 'alert') -> dict:
        """生成指定疲劳等级的配置"""
        params = self.fatigue_levels[level]
        return {
            'level': level,
            'blink_rate': np.random.uniform(*params['blink_rate']),
            'closing_ms': np.random.uniform(*params['closing_duration']),
            'closed_ms': np.random.uniform(*params['closed_duration']),
            'reopen_ms': np.random.uniform(*params['reopen_duration']),
            'amplitude': np.random.uniform(*params['amplitude']),
            'avatar': np.random.choice(self.parameters['avatar_models']),
            'camera': np.random.choice(self.parameters['camera_positions']),
            'lighting': np.random.choice(self.parameters['lighting']),
            'noise_db': np.random.choice(self.parameters['noise_levels']),
        }
    
    def generate_dataset(self, n_samples: int = 1000) -> list:
        """生成合成数据集配置"""
        levels = ['alert'] * 40 + ['low_vigilant'] * 25 + ['drowsy'] * 25 + ['microsleep'] * 10
        configs = []
        for i in range(n_samples):
            level = np.random.choice(levels)
            configs.append(self.generate_config(level))
        return configs


# 测试
if __name__ == "__main__":
    gen = SyntheticBlinkGenerator()
    
    print("合成数据生成配置示例:")
    for level in ['alert', 'low_vigilant', 'drowsy', 'microsleep']:
        config = gen.generate_config(level)
        print(f"\n{level}:")
        print(f"  眨眼频率: {config['blink_rate']:.1f} 次/分")
        print(f"  闭眼时长: {config['closed_ms']:.0f} ms")
        print(f"  重新睁眼: {config['reopen_ms']:.0f} ms")
        print(f"  振幅: {config['amplitude']:.1f}°")
    
    dataset = gen.generate_dataset(1000)
    print(f"\n数据集: {len(dataset)} 样本")
    level_counts = {}
    for c in dataset:
        level_counts[c['level']] = level_counts.get(c['level'], 0) + 1
    for level, count in sorted(level_counts.items()):
        print(f"  {level}: {count} ({count/len(dataset):.0%})")

# IMS 升级路线
upgrade_plan = {
    "current": {
        "metric": "EAR",
        "issues": ["头部旋转失真", "眼镜遮挡敏感", "个体差异大"],
        "accuracy": "~82%"
    },
    "upgrade_to_ELA": {
        "metric": "ELA (Eyelid Angle)",
        "benefits": [
            "头部旋转 >45° 仍稳定",
            "3D几何约束更鲁棒",
            "可直接计算眨眼时间特征"
        ],
        "accuracy": "~90%",
        "implementation": "MediaPipe Face Mesh → ELA → 眨眼特征 → 分类",
        "cost": "零额外硬件（仅算法升级）"
    },
    "synthetic_augmentation": {
        "benefit": "无需危险疲劳驾驶采集",
        "tool": "Blender 3D + rigged avatar",
        "volume": "1000+合成样本/天",
        "quality": "可控光照/视角/噪声/疲劳等级"
    }
}

# 合成数据增强方案
synthetic_training = {
    "problem": "真实疲劳数据稀缺且采集危险",
    "solution": "ELA驱动Blender生成合成数据",
    "pipeline": [
        "1. 采集少量真实疲劳ELA信号",
        "2. 提取眨眼参数分布",
        "3. 用参数驱动Blender avatar动画",
        "4. 多视角/光照/噪声渲染",
        "5. 生成标注数据",
        "6. 真实+合成混合训练"
    ],
    "cost_reduction": "数据采集成本降低80%+",
    "diversity": "4个avatar × 5视角 × 4光照 × 4噪声 = 320种配置",
    "quality_control": "可精确控制疲劳等级标注"
}