"""
免校准EEG驾驶员嗜睡检测框架

核心组件：
1. 预训练模型（源域数据训练）
2. Memory Bank（管理流式目标数据）
3. BN层自适应（仅更新BN层参数）
4. Prototype Learning（原型学习增强鲁棒性）
"""

import torch
import torch.nn as nn
import numpy as np
from collections import deque

class DrowsinessTTAFramework:
    """
    在线TTA驾驶员嗜睡检测框架
    
    论文核心方法复现
    
    工作流程：
    1. 预训练阶段：用多源驾驶员EEG训练基础模型
    2. 推理阶段：新驾驶员实时EEG流入
    3. 自适应阶段：通过Memory Bank + BN更新 + 原型学习
    """
    
    def __init__(self, model, n_classes=2, memory_size=64):
        self.model = model  # 预训练的EEG分类模型
        self.n_classes = n_classes
        self.memory_bank = MemoryBank(
            max_size=memory_size,
            n_classes=n_classes
        )
        self.prototype_layer = PrototypeLayer(
            n_classes=n_classes,
            feat_dim=128
        )
        
        # 冻结所有层，只保留BN层可训练
        self._freeze_non_bn_layers()
    
    def _freeze_non_bn_layers(self):
        """冻结BN层以外的所有参数"""
        for name, param in self.model.named_parameters():
            if 'bn' in name.lower():
                param.requires_grad = True  # BN层可更新
            else:
                param.requires_grad = False  # 其他层冻结
        
        # 但保留预训练BN统计量
        for m in self.model.modules():
            if isinstance(m, nn.BatchNorm1d):
                m.track_running_stats = True  # 保留预训练统计量
                # 仅更新affine参数(weight, bias)

class MemoryBank:
    """
    动态Memory Bank
    
    功能：
    1. 存储流式EEG片段
    2. 基于可靠性（negative energy score）筛选
    3. 基于持久时间管理样本生命周期
    """
    
    def __init__(self, max_size=64, n_classes=2, 
                 energy_threshold=0.5, max_persistence=50):
        self.max_size = max_size
        self.n_classes = n_classes
        self.energy_threshold = energy_threshold
        self.max_persistence = max_persistence
        
        # 存储结构: {class_id: deque([(sample, age)])}
        self.storage = {
            i: deque(maxlen=max_size // n_classes)
            for i in range(n_classes)
        }
    
    def add_sample(self, sample, prediction, energy_score):
        """
        添加新样本到Memory Bank
        
        Args:
            sample: EEG特征tensor, shape=(C, T)
            prediction: 模型预测的类别
            energy_score: negative energy score（可靠性指标）
        
        判断逻辑：
        - energy_score < threshold → 高可靠性 → 接受
        - persistence超过max_persistence → 移除（过时数据）
        """
        # 可靠性判断
        if energy_score < self.energy_threshold:
            # 高可靠性样本，加入Memory Bank
            self.storage[prediction].append({
                'sample': sample,
                'age': 0,
                'energy': energy_score
            })
        
        # 老化管理
        self._age_and_remove()
    
    def _age_and_remove(self):
        """增加所有样本年龄，移除过时样本"""
        for class_id in self.storage:
            to_remove = []
            for i, item in enumerate(self.storage[class_id]):
                item['age'] += 1
                if item['age'] > self.max_persistence:
                    to_remove.append(i)
            
            # 从后往前删除
            for i in reversed(to_remove):
                del self.storage[class_id][i]
    
    def get_reliable_samples(self):
        """获取当前Memory Bank中所有可靠样本"""
        samples = []
        labels = []
        for class_id, items in self.storage.items():
            for item in items:
                samples.append(item['sample'])
                labels.append(class_id)
        return torch.stack(samples), torch.tensor(labels)

def negative_energy_score(logits, temperature=1.0):
    """
    Negative Energy Score
    
    基于能量函数评估样本可靠性：
    - 能量低 → 样本在模型已知分布内 → 可靠
    - 能量高 → 样本在分布外 → 不可靠
    
    公式: E(x) = -T * log(sum(exp(logits / T)))
    
    Args:
        logits: shape=(batch, n_classes), 模型输出
        temperature: 温度参数
        
    Returns:
        energy: shape=(batch,), negative energy scores
    """
    # 能量函数
    energy = -temperature * torch.log(
        torch.sum(torch.exp(logits / temperature), dim=1)
    )
    return energy


# 使用示例
if __name__ == "__main__":
    # 模拟模型输出
    logits = torch.tensor([
        [2.5, 0.3],   # 高置信度（清醒）
        [0.1, 2.8],   # 高置信度（嗜睡）
        [0.5, 0.6],   # 低置信度（不确定）
    ])
    
    scores = negative_energy_score(logits, T=1.0)
    print(f"Energy scores: {scores}")
    # 低score → 高可靠性
    # sample 0 & 1 会被接受, sample 2 会被拒绝

class PrototypeLayer(nn.Module):
    """
    原型学习层
    
    为每个类别维护一个原型向量（prototype vector）
    分类基于样本到原型的距离，而非softmax
    
    优势：
    1. 对分布漂移鲁棒（距离度量比概率更稳定）
    2. 原型可随Memory Bank更新
    """
    
    def __init__(self, n_classes=2, feat_dim=128):
        super().__init__()
        self.n_classes = n_classes
        self.feat_dim = feat_dim
        
        # 原型向量
        self.prototypes = nn.Parameter(
            torch.randn(n_classes, feat_dim)
        )
    
    def update_prototypes(self, features, labels):
        """
        基于Memory Bank更新原型
        
        Args:
            features: shape=(n_samples, feat_dim), Memory Bank中的特征
            labels: shape=(n_samples,), 对应标签
        """
        for c in range(self.n_classes):
            mask = (labels == c)
            if mask.sum() > 0:
                class_features = features[mask]
                # 原型 = 类内特征均值
                new_prototype = class_features.mean(dim=0)
                # 指数移动平均更新
                momentum = 0.9
                self.prototypes.data[c] = (
                    momentum * self.prototypes.data[c] +
                    (1 - momentum) * new_prototype
                )
    
    def forward(self, features):
        """
        计算样本到各原型的距离
        
        Args:
            features: shape=(batch, feat_dim)
            
        Returns:
            distances: shape=(batch, n_classes)
        """
        # 欧氏距离
        dist = torch.cdist(features, self.prototypes, p=2)
        # 转换为概率（距离越近概率越大）
        probs = torch.softmax(-dist / 0.1, dim=1)
        return probs

class OnlineDrowsinessDetector:
    """
    完整的在线嗜睡检测器
    
    推理流程：
    1. 接收实时EEG片段
    2. 模型前向推理
    3. 计算negative energy score
    4. 可靠样本加入Memory Bank
    5. 定期用Memory Bank更新BN参数和原型
    6. 输出嗜睡概率
    """
    
    def __init__(self, pretrained_model, update_interval=10):
        self.tta_framework = DrowsinessTTAFramework(pretrained_model)
        self.update_interval = update_interval
        self.sample_count = 0
        
    def predict(self, eeg_segment):
        """
        实时嗜睡检测
        
        Args:
            eeg_segment: shape=(1, n_channels, n_samples)
                         单个EEG片段（通常2-4秒）
        
        Returns:
            drowsiness_prob: float, 嗜睡概率
            is_reliable: bool, 样本是否可靠
        """
        # Step 1: 前向推理
        with torch.no_grad():
            features = self.tta_framework.model.extract_features(eeg_segment)
            logits = self.tta_framework.model.classify(features)
        
        # Step 2: 计算能量分数
        energy = negative_energy_score(logits)
        is_reliable = energy.item() < 0.5
        
        # Step 3: 可靠样本加入Memory Bank
        pred_class = logits.argmax(dim=1).item()
        self.tta_framework.memory_bank.add_sample(
            sample=features.squeeze(0),
            prediction=pred_class,
            energy_score=energy.item()
        )
        
        # Step 4: 定期自适应更新
        self.sample_count += 1
        if self.sample_count % self.update_interval == 0:
            self._adapt_model()
        
        # Step 5: 输出嗜睡概率
        drowsiness_prob = torch.softmax(logits, dim=1)[0, 1].item()
        
        return drowsiness_prob, is_reliable
    
    def _adapt_model(self):
        """使用Memory Bank更新模型"""
        samples, labels = self.tta_framework.memory_bank.get_reliable_samples()
        
        if len(samples) < 4:  # 至少需要4个样本
            return
        
        # 前向传播
        features = self.tta_framework.model.extract_features(samples)
        
        # 更新BN层affine参数
        loss = self._compute_loss(features, labels)
        loss.backward()  # 只更新BN层（其他已冻结）
        
        # 更新原型
        self.tta_framework.prototype_layer.update_prototypes(
            features.detach(), labels
        )
    
    def _compute_loss(self, features, labels):
        """计算自适应损失"""
        # 原型距离损失 + 熵最小化
        proto_probs = self.tta_framework.prototype_layer(features)
        proto_loss = nn.CrossEntropyLoss()(proto_probs, labels)
        
        # 熵最小化（鼓励自信预测）
        logits = self.tta_framework.model.classify(features)
        entropy = -torch.sum(
            torch.softmax(logits, dim=1) * 
            torch.log(torch.softmax(logits, dim=1) + 1e-10),
            dim=1
        ).mean()
        
        return proto_loss + 0.1 * entropy

"""
座舱EEG采集方案对比

方案1: 方向盘电极（最接近量产）
- 电极位置: 手握方向盘时接触
- 通道数: 2-4 (有限)
- 信号质量: 中等（运动伪影）
- 用户接受度: 高（无额外操作）

方案2: 头枕电极
- 电极位置: 头枕接触后脑
- 通道数: 4-8
- 信号质量: 好（接触稳定）
- 用户接受度: 中等

方案3: 可穿戴EEG头环
- 电极位置: 前额+耳垂
- 通道数: 8-16
- 信号质量: 好
- 用户接受度: 低（需佩戴额外设备）
"""

# 方案1最实用，但通道数有限
# TTA的价值：即使2-4通道也能达到81.73% F1
# 传统方法在2通道时F1可能<50%

STEERING_WHEEL_EEG_CONFIG = {
    'channels': ['Fp1', 'Fp2', 'A1', 'A2'],  # 前额+参考
    'sampling_rate': 250,  # Hz
    'resolution': 16,  # bit
    'contact_type': 'dry_electrode',  # 干电极
    'max_impedance': 100,  # kOhm
    'update_rate': 2,  # Hz (每0.5s输出一次)
}

"""
座舱EEG-TTA改进方案

改进1: 增加rPPG作为辅助模态
- rPPG提供心率变异性，与EEG互补
- 无需额外传感器，复用DMS摄像头
"""
class MultiModalTTA(DrowsinessTTAFramework):
    def __init__(self, eeg_model, rppg_model):
        super().__init__(eeg_model)
        self.rppg_model = rppg_model
        
    def predict(self, eeg_segment, rppg_signal):
        """EEG+rPPG融合嗜睡检测"""
        # EEG分支
        eeg_prob, reliable = super().predict(eeg_segment)
        
        # rPPG分支
        hr = self._extract_heart_rate(rppg_signal)
        hrv = self._compute_hrv(hr)
        rppg_prob = self.rppg_model.predict(hrv)
        
        # 融合
        if reliable:
            # EEG可靠时权重高
            fused = 0.7 * eeg_prob + 0.3 * rppg_prob
        else:
            # EEG不可靠时依赖rPPG
            fused = 0.3 * eeg_prob + 0.7 * rppg_prob
        
        return fused


"""
改进2: 持续学习防止灾难性遗忘
"""
class ContinualTTA(DrowsinessTTAFramework):
    """持续TTA with EWC正则化"""
    
    def __init__(self, model, ewc_lambda=0.01):
        super().__init__(model)
        self.ewc_lambda = ewc_lambda
        self.fisher_matrix = None  # Fisher信息矩阵
    
    def compute_fisher(self, data_loader):
        """计算Fisher信息矩阵（初始化时）"""
        self.fisher_matrix = {}
        for name, param in self.model.named_parameters():
            if param.requires_grad:
                self.fisher_matrix[name] = torch.zeros_like(param)
        
        for x, y in data_loader:
            self.model.zero_grad()
            output = self.model(x)
            loss = nn.CrossEntropyLoss()(output, y)
            loss.backward()
            for name, param in self.model.named_parameters():
                if param.requires_grad:
                    self.fisher_matrix[name] += param.grad.data ** 2
    
    def _adapt_model(self):
        """带EWC正则化的自适应"""
        samples, labels = self.tta_framework.memory_bank.get_reliable_samples()
        features = self.model.extract_features(samples)
        
        loss = self._compute_loss(features, labels)
        
        # EWC正则化
        if self.fisher_matrix:
            ewc_loss = 0
            for name, param in self.model.named_parameters():
                if param.requires_grad and name in self.fisher_matrix:
                    ewc_loss += (self.fisher_matrix[name] * 
                                (param - param.data) ** 2).sum()
            loss += self.ewc_lambda * ewc_loss
        
        loss.backward()