import cv2
import numpy as np
from typing import List, Tuple
import mediapipe as mp

class RailwayDrowsinessPreprocessor:
    """
    铁路驾驶员疲劳检测数据预处理
    
    论文方法复现：
    1. 视频帧提取与降采样
    2. 面部关键点检测（12个关键点）
    3. EAR（眼睑开度比）计算
    4. MAR（嘴巴开度比）计算
    5. 数据增强
    """
    
    def __init__(self):
        # 使用 MediaPipe + cvzone 进行面部关键点检测
        self.mp_face_mesh = mp.solutions.face_mesh
        self.face_mesh = self.mp_face_mesh.FaceMesh(
            max_num_faces=1,
            refine_landmarks=True,
            min_detection_confidence=0.5,
            min_tracking_confidence=0.5
        )
        
        # 眼睛关键点索引（MediaPipe 468点模型）
        self.LEFT_EYE = [33, 160, 158, 133, 153, 144]  # 6点
        self.RIGHT_EYE = [362, 385, 387, 263, 373, 380]  # 6点
        self.MOUTH = [13, 14, 78, 308]  # 4点
    
    def extract_frames(
        self,
        video_path: str,
        target_fps: int = 30,
        sample_ratio: int = 6
    ) -> List[np.ndarray]:
        """
        帧提取与降采样
        
        论文方法：每6帧取1帧（2,913/17,476）
        
        Args:
            video_path: 视频文件路径
            target_fps: 目标帧率
            sample_ratio: 采样比率
            
        Returns:
            frames: 采样后的帧列表
        """
        cap = cv2.VideoCapture(video_path)
        frames = []
        frame_idx = 0
        
        while True:
            ret, frame = cap.read()
            if not ret:
                break
            
            # 每6帧取1帧
            if frame_idx % sample_ratio == 0:
                frames.append(frame)
            
            frame_idx += 1
        
        cap.release()
        return frames
    
    def detect_facial_landmarks(
        self,
        frame: np.ndarray
    ) -> Tuple[np.ndarray, bool]:
        """
        检测12个面部关键点
        
        论文方法：使用 cvzone/MediaPipe 检测
        - 左眼4点 + 右眼4点 + 嘴巴4点 = 12点
        """
        rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        results = self.face_mesh.process(rgb_frame)
        
        if results.multi_face_landmarks:
            landmarks = results.multi_face_landmarks[0].landmark
            points = []
            for idx in self.LEFT_EYE[:4] + self.RIGHT_EYE[:4] + self.MOUTH:
                points.append([landmarks[idx].x, landmarks[idx].y])
            return np.array(points), True
        
        return np.array([]), False
    
    def calculate_ear(self, eye_points: np.ndarray) -> float:
        """
        计算 Eye Aspect Ratio (EAR)
        
        EAR = (|p2-p6| + |p3-p5|) / (2 * |p1-p4|)
        
        论文阈值：EAR < 18 判定为闭眼
        注：论文使用的阈值与标准 EAR 公式不同，
        可能是 cvzone 的 Ratio 实现
        
        Args:
            eye_points: 4个眼角点 (x, y)
            
        Returns:
            ear: 眼睑开度比值
        """
        # 简化为4点版本的 EAR
        # p1=左角, p2=上, p3=右角, p4=下
        horizontal = np.linalg.norm(eye_points[0] - eye_points[2])
        vertical = np.linalg.norm(eye_points[1] - eye_points[3])
        
        if horizontal > 0:
            ratio = vertical / horizontal
        else:
            ratio = 0
        
        return ratio
    
    def calculate_mar(self, mouth_points: np.ndarray) -> float:
        """
        计算 Mouth Aspect Ratio (MAR)
        
        论文阈值：MAR > 35 判定为打哈欠
        """
        # 嘴角水平距离
        horizontal = np.linalg.norm(mouth_points[0] - mouth_points[2])
        # 嘴唇垂直距离
        vertical = np.linalg.norm(mouth_points[1] - mouth_points[3])
        
        if horizontal > 0:
            ratio = vertical / horizontal * 100  # 论文使用放大系数
        else:
            ratio = 0
        
        return ratio
    
    def augment_frame(
        self,
        frame: np.ndarray,
        rotation_range: Tuple[float, float] = (-15, 15),
        brightness_range: Tuple[float, float] = (-0.2, 0.2)
    ) -> np.ndarray:
        """
        数据增强
        
        论文方法：
        - 旋转：-15° 到 +15°
        - 亮度：-20% 到 +20%
        
        使用 Roboflow 进行数据增强
        """
        h, w = frame.shape[:2]
        
        # 随机旋转
        angle = np.random.uniform(*rotation_range)
        M = cv2.getRotationMatrix2D((w/2, h/2), angle, 1)
        frame = cv2.warpAffine(frame, M, (w, h))
        
        # 随机亮度调整
        brightness = np.random.uniform(*brightness_range)
        frame = np.clip(frame * (1 + brightness), 0, 255).astype(np.uint8)
        
        return frame


# ============ YOLOv8 分类模型 ============

class RailwayDrowsinessClassifier:
    """
    基于 YOLOv8 的铁路驾驶员疲劳分类器
    
    论文对比三个变体：
    | 模型 | 参数量 | 最佳Epoch | 损失值 |
    |------|--------|----------|--------|
    | YOLOv8n-cls | ~2.7M | 145 | 0.06787 |
    | YOLOv8s-cls | ~11M | 100 | 0.04662 |
    | YOLOv8m-cls | ~26M | 64 | 0.05495 |
    
    最终选择 YOLOv8n-cls（轻量级最优）
    
    训练配置：
    - Image size: 224 × 224 × 3
    - Optimizer: AdamW
    - Learning rate: 0.01
    - Batch size: 1024 (n) / 256 (s)
    """
    
    def __init__(self, model_variant: str = 'n'):
        """
        Args:
            model_variant: 'n' (nano), 's' (small), 'm' (medium)
        """
        self.variant = model_variant
        self.model_name = f"yolov8{model_variant}-cls"
        
    def get_config(self) -> dict:
        """返回论文训练配置"""
        configs = {
            'n': {
                'epochs': 500,
                'batch_size': 1024,
                'learning_rate': 0.01,
                'image_size': 224,
                'optimizer': 'AdamW',
                'best_epoch': 145,
                'final_loss': 0.06787,
                'params': 2_700_000  # ~2.7M
            },
            's': {
                'epochs': 100,
                'batch_size': 256,
                'learning_rate': 0.01,
                'image_size': 224,
                'optimizer': 'AdamW',
                'best_epoch': 100,
                'final_loss': 0.04662,
                'params': 11_000_000
            },
            'm': {
                'epochs': 500,
                'batch_size': 1024,
                'learning_rate': 0.01,
                'image_size': 224,
                'optimizer': 'AdamW',
                'best_epoch': 64,
                'final_loss': 0.05495,
                'params': 26_000_000
            }
        }
        return configs[self.variant]
    
    def get_training_script(self) -> str:
        """生成训练脚本"""
        config = self.get_config()
        return f"""
# 铁路驾驶员疲劳检测 - YOLOv8{self.variant}-cls 训练脚本
# 论文复现

from ultralytics import YOLO

# 加载预训练模型
model = YOLO('yolov8{self.variant}-cls.pt')

# 训练
results = model.train(
    data='./railway_drowsiness_dataset/',  # 数据集路径
    epochs={config['epochs']},
    batch={config['batch_size']},
    lr0={config['learning_rate']},
    imgsz={config['image_size']},
    optimizer='{config['optimizer']}',
    # 论文使用 Roboflow 进行数据增强
    hsv_h=0.0,  # 关闭色彩增强（铁路场景色彩固定）
    hsv_s=0.0,
    hsv_v=0.2,  # 亮度 ±20%
    degrees=15.0,  # 旋转 ±15°
    translate=0.1,
    scale=0.0,
    fliplr=0.0,  # 水平翻转（人脸不翻转）
    flipud=0.0,
    patience=50,  # 早停
    device='cuda:0',
    workers=8
)

# 验证
metrics = model.val()
print(f"Accuracy: {{metrics.top1:.4f}}")
print(f"Precision: {{metrics.precision:.4f}}")
print(f"Recall: {{metrics.recall:.4f}}")
"""


# ============ 实时检测系统 ============

class RealTimeDrowsinessSystem:
    """
    实时铁路驾驶员疲劳检测系统
    
    部署硬件：NVIDIA A100 40GB (论文使用)
    实际量产可用更轻量设备
    
    检测逻辑：
    1. 实时视频流 → 面部检测
    2. EAR < 阈值 + 持续时间 > 833ms → 困倦
    3. MAR > 阈值 → 打哈欠
    4. 触发视觉/声音告警
    """
    
    def __init__(self, config: dict = None):
        config = config or {}
        self.ear_threshold = config.get('ear_threshold', 0.18)
        self.mar_threshold = config.get('mar_threshold', 0.35)
        self.eye_closure_duration_threshold = 0.833  # 833ms
        
        self.eye_closure_start = None
        self.drowsy_state = False
        self.alert_count = 0
        
    def process_frame(
        self,
        frame: np.ndarray,
        eye_landmarks: np.ndarray,
        mouth_landmarks: np.ndarray,
        timestamp: float
    ) -> dict:
        """
        处理单帧
        
        Args:
            frame: 视频帧
            eye_landmarks: 眼睛关键点
            mouth_landmarks: 嘴巴关键点
            timestamp: 时间戳（秒）
            
        Returns:
            result: 检测结果
        """
        ear = self._calc_ear(eye_landmarks)
        mar = self._calc_mar(mouth_landmarks)
        
        result = {
            'ear': ear,
            'mar': mar,
            'is_drowsy': False,
            'is_yawning': False,
            'alert_level': 0,
            'timestamp': timestamp
        }
        
        # 检测闭眼
        if ear < self.ear_threshold:
            if self.eye_closure_start is None:
                self.eye_closure_start = timestamp
            else:
                duration = timestamp - self.eye_closure_start
                if duration >= self.eye_closure_duration_threshold:
                    result['is_drowsy'] = True
                    result['alert_level'] = 1
                    self.alert_count += 1
        else:
            self.eye_closure_start = None
        
        # 检测打哈欠
        if mar > self.mar_threshold:
            result['is_yawning'] = True
            if result['alert_level'] < 1:
                result['alert_level'] = 1
        
        # 连续困倦 → 升级告警
        if self.alert_count >= 3:
            result['alert_level'] = 2
        
        return result
    
    def _calc_ear(self, points: np.ndarray) -> float:
        v1 = np.linalg.norm(points[1] - points[5])
        v2 = np.linalg.norm(points[2] - points[4])
        h = np.linalg.norm(points[0] - points[3])
        return (v1 + v2) / (2.0 * h) if h > 0 else 0
    
    def _calc_mar(self, points: np.ndarray) -> float:
        v = np.linalg.norm(points[1] - points[3])
        h = np.linalg.norm(points[0] - points[2])
        return v / h if h > 0 else 0


if __name__ == "__main__":
    # 预处理器测试
    preprocessor = RailwayDrowsinessPreprocessor()
    print("预处理模块初始化完成")
    
    # 模型配置
    classifier = RailwayDrowsinessClassifier('n')
    config = classifier.get_config()
    print(f"\n模型: YOLOv8n-cls")
    print(f"参数量: {config['params']:,}")
    print(f"训练轮数: {config['epochs']}")
    print(f"最佳轮次: {config['best_epoch']}")
    print(f"最终损失: {config['final_loss']}")
    
    # 训练脚本
    print("\n训练脚本:")
    print(classifier.get_training_script())
    
    # 实时系统
    system = RealTimeDrowsinessSystem()
    print("\n实时检测系统初始化完成")
    print(f"闭眼阈值: EAR < {system.ear_threshold}")
    print(f"闭眼时长阈值: {system.eye_closure_duration_threshold * 1000:.0f}ms")
    print(f"打哈欠阈值: MAR > {system.mar_threshold}")
    
    print("\n=== 论文报告结果 ===")
    print("总准确率: 96.8%")
    print("精确率: 97.28%")
    print("召回率: 97.46%")
    print("\n各时段准确率:")
    print("  早上: ~98.0%")
    print("  下午: ~95.4% (日光影响)")
    print("  夜间: ~97.0%")