import tensorflow as tf
import numpy as np
from tensorflow.keras.applications import MobileNetV2
from tensorflow.keras.layers import Dense, GlobalAveragePooling2D
from tensorflow.keras.models import Model

def build_branch(input_size=(96, 96, 3), num_classes=2, 
                 name='branch'):
    """
    构建单分支 MobileNetV2 分类器
    
    Args:
        input_size: 输入尺寸
        num_classes: 类别数
        name: 分支名称
    """
    base = MobileNetV2(
        input_shape=input_size,
        include_top=False,
        weights='imagenet',
    )
    base.trainable = False  # 冻结骨干
    
    x = base.output
    x = GlobalAveragePooling2D()(x)
    x = Dense(128, activation='relu')(x)
    outputs = Dense(num_classes, activation='softmax')(x)
    
    return Model(base.input, outputs, name=name)

# 构建双分支模型
eye_model = build_branch(name='eye_branch')  # Open/Closed
mouth_model = build_branch(name='mouth_branch')  # Yawn/No_yawn

# 训练参数
EYE_CLASSES = ['Open', 'Closed']
MOUTH_CLASSES = ['Yawn', 'No_yawn']

# 数据增强
datagen = tf.keras.preprocessing.image.ImageDataGenerator(
    rotation_range=10,
    width_shift_range=0.1,
    height_shift_range=0.1,
    horizontal_flip=True,
    rescale=1./255,
)

# 编译训练
eye_model.compile(
    optimizer='adam',
    loss='sparse_categorical_crossentropy',
    metrics=['accuracy'],
)

# 训练（伪代码）
# eye_model.fit(eye_train_gen, epochs=20, validation_data=eye_val_gen)

from collections import deque
from enum import Enum
import time

class DrowsinessState(Enum):
    ALERT = 'alert'
    DROWSY = 'drowsy'
    MICRO_SLEEP = 'micro_sleep'
    EXCESSIVE_YAWN = 'excessive_yawn'

class DualModelDecisionEngine:
    """
    双模型时间决策引擎
    
    将帧级分类结果转换为时序告警：
    - 桌面版：基于持续时间的双级滞回告警
    - 边缘版：基于帧计数的告警
    """
    
    def __init__(self, fps: float, platform: str = 'desktop'):
        """
        Args:
            fps: 视频帧率
            platform: 'desktop' 或 'edge'
        """
        self.fps = fps
        self.platform = platform
        
        # 滑动窗口
        self.window_size = int(fps * 60)  # 60秒窗口
        self.eye_states = deque(maxlen=self.window_size)
        self.mouth_states = deque(maxlen=self.window_size)
        
        # 告警参数
        if platform == 'desktop':
            # 桌面：基于持续时间（秒）
            self.eye_closed_threshold_sec = 1.5
            self.micro_sleep_threshold_sec = 2.0
            self.yawn_threshold_sec = 3.0
            self.yawn_rate_threshold = 3  # 次/分钟
        else:
            # 边缘：基于帧计数
            self.eye_closed_threshold_frames = int(1.5 * fps)
            self.micro_sleep_threshold_frames = int(2.0 * fps)
            self.yawn_threshold_frames = int(3.0 * fps)
            self.yawn_rate_threshold = 3
        
        # 状态跟踪
        self.eye_closed_start = None
        self.yawn_start = None
        self.yawn_count = 0
        self.alerts = []
    
    def update(self, eye_pred: str, mouth_pred: str) -> dict:
        """
        更新一帧的预测结果
        
        Args:
            eye_pred: 'Open' 或 'Closed'
            mouth_pred: 'Yawn' 或 'No_yawn'
        """
        timestamp = time.time()
        self.eye_states.append((timestamp, eye_pred))
        self.mouth_states.append((timestamp, mouth_pred))
        
        # 检测闭眼持续时间
        if eye_pred == 'Closed':
            if self.eye_closed_start is None:
                self.eye_closed_start = timestamp
            closed_duration = timestamp - self.eye_closed_start
            if closed_duration >= self.micro_sleep_threshold_sec:
                self._trigger_alert('micro_sleep', closed_duration)
        else:
            if self.eye_closed_start is not None:
                duration = timestamp - self.eye_closed_start
                if duration >= self.eye_closed_threshold_sec:
                    self._trigger_alert('drowsy', duration)
                self.eye_closed_start = None
        
        # 检测打哈欠
        if mouth_pred == 'Yawn':
            if self.yawn_start is None:
                self.yawn_start = timestamp
            yawn_duration = timestamp - self.yawn_start
            if yawn_duration >= self.yawn_threshold_sec:
                self._trigger_alert('excessive_yawn', yawn_duration)
        else:
            if self.yawn_start is not None:
                self.yawn_count += 1
                self.yawn_start = None
        
        # 检查打哈欠频率
        recent_yawns = self._count_recent_yawns(60)  # 最近60秒
        if recent_yawns >= self.yawn_rate_threshold:
            self._trigger_alert('excessive_yawn_rate', recent_yawns)
        
        return self.get_current_state()
    
    def _count_recent_yawns(self, seconds: int) -> int:
        """统计最近N秒内的打哈欠次数"""
        now = time.time()
        cutoff = now - seconds
        # 简化：统计闭眼-睁眼转换次数
        count = 0
        for i in range(1, len(self.eye_states)):
            if self.eye_states[i][1] == 'Open' and \
               self.eye_states[i-1][1] == 'Closed':
                if self.eye_states[i][0] > cutoff:
                    count += 1
        return count
    
    def _trigger_alert(self, alert_type: str, value: float):
        """触发告警"""
        self.alerts.append({
            'timestamp': time.time(),
            'type': alert_type,
            'value': value,
        })
    
    def get_current_state(self) -> dict:
        """获取当前状态"""
        if not self.alerts:
            return {'state': DrowsinessState.ALERT, 'alerts': []}
        
        latest = self.alerts[-1]
        if latest['type'] == 'micro_sleep':
            state = DrowsinessState.MICRO_SLEEP
        elif latest['type'] in ('excessive_yawn', 'excessive_yawn_rate'):
            state = DrowsinessState.EXCESSIVE_YAWN
        else:
            state = DrowsinessState.DROWSY
        
        return {
            'state': state,
            'alerts': self.alerts[-5:],  # 最近5条
        }


# 测试
if __name__ == "__main__":
    import random
    
    engine = DualModelDecisionEngine(fps=30, platform='desktop')
    
    # 模拟30秒视频（900帧）
    print("=== 模拟正常驾驶（30秒）===")
    for i in range(900):
        eye = 'Open' if random.random() > 0.02 else 'Closed'
        mouth = 'No_yawn' if random.random() > 0.01 else 'Yawn'
        result = engine.update(eye, mouth)
    
    print(f"状态: {result['state']}")
    print(f"告警数: {len(result['alerts'])}")
    
    # 模拟疲劳驾驶（微睡眠+打哈欠）
    print("\n=== 模拟疲劳驾驶（30秒）===")
    for i in range(900):
        if 200 <= i <= 260:  # 第7-9秒：闭眼2秒
            eye = 'Closed'
        elif 400 <= i <= 420:  # 第13-14秒：闭眼
            eye = 'Closed'
        else:
            eye = 'Open' if random.random() > 0.05 else 'Closed'
        
        if 300 <= i <= 390:  # 第10-13秒：打哈欠
            mouth = 'Yawn'
        else:
            mouth = 'No_yawn'
        
        result = engine.update(eye, mouth)
    
    print(f"状态: {result['state']}")
    print(f"告警数: {len(result['alerts'])}")
    for alert in result['alerts']:
        print(f"  {alert['type']}: {alert['value']:.1f}")

def convert_to_tflite(keras_model, output_path: str):
    """
    将 Keras H5 模型转换为 TFLite 量化模型
    
    论文报告结果：
    - 模型大小：8.99 MB → 2.39 MB (↓73.4%)
    - 推理延迟：~25 ms → ~1.9 ms (↓13x)
    """
    converter = tf.lite.TFLiteConverter.from_keras_model(keras_model)
    
    # 全整数量化
    converter.optimizations = [tf.lite.Optimize.DEFAULT]
    converter.target_spec.supported_types = [tf.int8]
    converter.target_spec.supported_ops = [
        tf.lite.OpsSet.TFLITE_BUILTINS_INT8,
    ]
    
    # 代表性数据集（用于校准）
    def representative_dataset():
        for _ in range(100):
            data = np.random.rand(1, 96, 96, 3).astype(np.float32)
            yield [data]
    
    converter.representative_dataset = representative_dataset
    
    tflite_model = converter.convert()
    
    with open(output_path, 'wb') as f:
        f.write(tflite_model)
    
    # 打印大小
    import os
    size_mb = os.path.getsize(output_path) / 1024 / 1024
    print(f"TFLite 模型: {output_path} ({size_mb:.2f} MB)")
    
    return output_path


def benchmark_tflite(model_path: str, num_runs: int = 100):
    """
    基准测试 TFLite 模型推理延迟
    """
    interpreter = tf.lite.Interpreter(model_path=model_path)
    interpreter.allocate_tensors()
    
    input_details = interpreter.get_input_details()[0]
    input_shape = input_details['shape']
    
    # 预热
    test_input = np.random.rand(*input_shape).astype(np.float32)
    interpreter.set_tensor(input_details['index'], test_input)
    interpreter.invoke()
    
    # 正式测试
    times = []
    for _ in range(num_runs):
        start = time.time()
        interpreter.set_tensor(input_details['index'], test_input)
        interpreter.invoke()
        times.append((time.time() - start) * 1000)
    
    avg_ms = np.mean(times)
    p95_ms = np.percentile(times, 95)
    
    print(f"平均延迟: {avg_ms:.2f} ms")
    print(f"P95 延迟: {p95_ms:.2f} ms")
    print(f"等效帧率: {1000/avg_ms:.1f} FPS")
    
    return {'avg_ms': avg_ms, 'p95_ms': p95_ms}