# 示意代码,不可直接运行(本站按论文"置信度驱动自适应窗口"思想改写的简化示意;原文模型为
# MobileNetV3-CBAM-BiLSTM + MC Dropout,论文未提供官方开源仓库链接)
import torch
import torch.nn as nn
import numpy as np
from dataclasses import dataclass

@dataclass
class WindowState:
    """时间窗口状态"""
    current_size: float = 30.0  # 秒
    min_size: float = 5.0
    max_size: float = 60.0
    confidence: float = 0.5
    trend: str = 'stable'  # 'expanding', 'shrinking', 'stable'


class ConfidenceAdaptiveWindow:
    """
    置信度驱动自适应时间窗口
    
    原理：
    - 高置信度→缩短窗口→降低延迟
    - 低置信度→扩大窗口→提升精度
    - 类似生物注意力机制：信号可靠时快速反应
    """
    
    def __init__(self, 
                 min_window: float = 5.0,
                 max_window: float = 60.0,
                 init_window: float = 30.0,
                 confidence_high: float = 0.85,
                 confidence_low: float = 0.50,
                 expand_rate: float = 1.5,
                 shrink_rate: float = 0.7):
        self.min_window = min_window
        self.max_window = max_window
        self.current_window = init_window
        self.conf_high = confidence_high
        self.conf_low = confidence_low
        self.expand_rate = expand_rate
        self.shrink_rate = shrink_rate
        self.history = []
    
    def update(self, confidence: float) -> float:
        """
        根据预测置信度调整窗口大小
        
        Args:
            confidence: 当前预测置信度 [0, 1]
            
        Returns:
            new_window: 调整后的窗口大小（秒）
        """
        self.history.append(confidence)
        
        # 移动平均平滑
        if len(self.history) > 5:
            avg_conf = np.mean(self.history[-5:])
        else:
            avg_conf = confidence
        
        if avg_conf > self.conf_high:
            # 高置信→缩短窗口（更快响应）
            self.current_window *= self.shrink_rate
            self.current_window = max(self.current_window, self.min_window)
        elif avg_conf < self.conf_low:
            # 低置信→扩大窗口（更准）
            self.current_window *= self.expand_rate
            self.current_window = min(self.current_window, self.max_window)
        # 中等置信保持不变
        
        return self.current_window


class FatigueDetector(nn.Module):
    """
    疲劳检测模型（轻量CNN+MC Dropout）
    
    输出：疲劳分类 + 置信度
    """
    
    def __init__(self, n_classes: int = 3):
        super().__init__()
        self.backbone = nn.Sequential(
            nn.Conv2d(3, 16, 3, stride=2, padding=1),
            nn.BatchNorm2d(16),
            nn.ReLU6(),
            nn.Conv2d(16, 32, 3, stride=2, padding=1),
            nn.BatchNorm2d(32),
            nn.ReLU6(),
            nn.Conv2d(32, 64, 3, stride=2, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU6(),
            nn.AdaptiveAvgPool2d(1),
            nn.Flatten(),
        )
        self.classifier = nn.Sequential(
            nn.Linear(64, 32),
            nn.ReLU(),
            nn.Dropout(0.2),  # MC Dropout
            nn.Linear(32, n_classes)
        )
    
    def forward(self, x, n_samples=1):
        feat = self.backbone(x)
        
        if n_samples > 1:
            logits_list = []
            for _ in range(n_samples):
                logits_list.append(self.classifier(feat))
            logits_stack = torch.stack(logits_list)
            mean_logits = logits_stack.mean(dim=0)
            probs = torch.softmax(mean_logits, dim=-1)
            confidence = probs.max(dim=-1)[0]
            # 认知不确定性
            uncertainty = 1 - confidence
        else:
            mean_logits = self.classifier(feat)
            probs = torch.softmax(mean_logits, dim=-1)
            confidence = probs.max(dim=-1)[0]
            uncertainty = 1 - confidence
        
        return mean_logits, confidence, uncertainty


class AdaptiveFatigueSystem:
    """
    完整自适应疲劳检测系统
    
    架构：
    1. 自适应窗口管理器
    2. 疲劳检测模型
    3. 置信度反馈循环
    """
    
    def __init__(self):
        self.window_manager = ConfidenceAdaptiveWindow()
        self.detector = FatigueDetector(n_classes=3)
        self.fatigue_history = []
    
    def process_frame(self, frame: torch.Tensor) -> dict:
        """
        处理一帧：自适应窗口→检测→反馈
        
        Args:
            frame: [B, 3, H, W] 单帧或序列
        """
        # 1. 当前窗口大小
        window_size = self.window_manager.current_window
        
        # 2. 模型推理（10次MC采样）
        with torch.no_grad():
            logits, confidence, uncertainty = \
                self.detector(frame, n_samples=10)
        
        # 3. 置信度反馈→调整窗口
        new_window = self.window_manager.update(confidence.item())
        
        # 4. 决策
        probs = torch.softmax(logits, dim=-1)
        pred = probs.argmax(dim=-1)
        
        # 5. 报警策略
        if confidence > 0.85:
            alert_level = 'high_confidence'
        elif confidence > 0.50:
            alert_level = 'medium_confidence'
        else:
            alert_level = 'low_confidence_degraded'
        
        return {
            'prediction': pred.item(),
            'confidence': confidence.item(),
            'uncertainty': uncertainty.item(),
            'window_size': window_size,
            'new_window': new_window,
            'alert_level': alert_level,
        }


# 测试
if __name__ == "__main__":
    system = AdaptiveFatigueSystem()
    
    # 模拟清醒→疲劳过程
    print("=== 自适应窗口疲劳检测 ===")
    print(f"初始窗口: {system.window_manager.current_window:.0f}s")
    
    # 模拟20帧（清醒）
    print("\n--- 清醒阶段 ---")
    for i in range(10):
        frame = torch.randn(1, 3, 96, 96)
        result = system.process_frame(frame)
        if i % 3 == 0:
            print(f"  帧{i}: pred={result['prediction']}, "
                  f"conf={result['confidence']:.2f}, "
                  f"window={result['new_window']:.0f}s, "
                  f"level={result['alert_level']}")
    
    # 模拟疲劳（置信度下降）
    print("\n--- 疲劳阶段 ---")
    for i in range(10):
        # 模拟疲劳信号（低质量帧）
        frame = torch.randn(1, 3, 96, 96) * 0.5
        result = system.process_frame(frame)
        if i % 2 == 0:
            print(f"  帧{i}: pred={result['prediction']}, "
                  f"conf={result['confidence']:.2f}, "
                  f"window={result['new_window']:.0f}s, "
                  f"level={result['alert_level']}")