import numpy as np
import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader

"""
EEG 基础模型跨被试泛化评估框架
参考 EEG-FM-Compass (arXiv:2601.17883)

用于评估 EEG 基础模型在驾驶员认知分心检测中的跨被试迁移能力
"""

class EEGDataset(Dataset):
    """EEG 数据集，支持 LOSO 交叉验证"""
    
    def __init__(self, eeg_data: np.ndarray, labels: np.ndarray, 
                 subject_ids: np.ndarray, task: str = "cognitive_load"):
        """
        Args:
            eeg_data: EEG 信号, shape=(N, C, T)  N=样本数, C=通道数, T=时间点
            labels: 标签 (0=正常, 1=认知分心)
            subject_ids: 被试ID, shape=(N,)
            task: 任务类型
        """
        self.data = torch.FloatTensor(eeg_data)
        self.labels = torch.LongTensor(labels)
        self.subject_ids = subject_ids
        self.task = task
        
    def __len__(self):
        return len(self.data)
    
    def __getitem__(self, idx):
        return {
            'eeg': self.data[idx],
            'label': self.labels[idx],
            'subject_id': self.subject_ids[idx]
        }


class EEGTransformerEncoder(nn.Module):
    """
    轻量级 EEG Transformer 编码器
    参考 LaBraM 架构，适配驾驶员认知分心检测
    
    输入: (B, C, T)  B=batch, C=通道数(如19), T=时间点(如256, 2s@128Hz)
    输出: (B, num_classes) 认知状态分类
    """
    
    def __init__(self, 
                 num_channels: int = 19,
                 seq_len: int = 256,
                 embed_dim: int = 256,
                 num_heads: int = 8,
                 num_layers: int = 6,
                 num_classes: int = 2,
                 patch_size: int = 16):
        super().__init__()
        self.num_channels = num_channels
        self.patch_size = patch_size
        self.num_patches = seq_len // patch_size
        
        # 通道嵌入
        self.channel_embed = nn.Embedding(num_channels, embed_dim)
        
        # 时间片 patch 嵌入
        self.patch_embed = nn.Conv1d(
            in_channels=1, 
            out_channels=embed_dim,
            kernel_size=patch_size,
            stride=patch_size
        )
        
        # 位置编码
        self.pos_embed = nn.Parameter(
            torch.randn(1, num_channels * self.num_patches, embed_dim) * 0.02
        )
        
        # Transformer 编码器
        encoder_layer = nn.TransformerEncoderLayer(
            d_model=embed_dim,
            nhead=num_heads,
            dim_feedforward=embed_dim * 4,
            dropout=0.1,
            batch_first=True,
            activation='gelu'
        )
        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers)
        
        # 分类头
        self.norm = nn.LayerNorm(embed_dim)
        self.classifier = nn.Linear(embed_dim, num_classes)
        
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Args:
            x: (B, C, T) EEG 信号
        Returns:
            logits: (B, num_classes)
        """
        B, C, T = x.shape
        
        # 每通道切 patch
        x = x.unsqueeze(2)  # (B, C, 1, T)
        x = x.view(B * C, 1, T)
        x = self.patch_embed(x)  # (B*C, embed_dim, num_patches)
        x = x.transpose(1, 2)  # (B*C, num_patches, embed_dim)
        
        # 加入通道嵌入
        ch_ids = torch.arange(C, device=x.device).unsqueeze(0).repeat(B, 1)  # (B, C)
        ch_emb = self.channel_embed(ch_ids.view(-1))  # (B*C, embed_dim)
        ch_emb = ch_emb.view(B, C, 1, -1).repeat(1, 1, self.num_patches, 1)
        ch_emb = ch_emb.view(B * C, self.num_patches, -1)
        x = x + ch_emb
        
        # 位置编码
        x = x + self.pos_embed
        
        # Transformer
        x = self.transformer(x)
        x = self.norm(x)
        
        # 池化 + 分类
        x = x.mean(dim=1)  # (B*C, embed_dim)
        x = x.view(B, C, -1).mean(dim=1)  # (B, embed_dim)
        logits = self.classifier(x)
        
        return logits


class ProgressiveUnfreezing:
    """
    渐进式解冻策略
    参考 Fonya et al. (arXiv:2606.23706)
    
    逐步解冻 Transformer 层，避免灾难性遗忘
    """
    
    def __init__(self, model: nn.Module, total_epochs: int = 50, 
                 unfreeze_schedule: list = None):
        self.model = model
        self.total_epochs = total_epochs
        # 默认：每10个epoch解冻一层
        self.schedule = unfreeze_schedule or list(range(0, total_epochs, 10))
        
    def apply_freeze(self, epoch: int):
        """根据当前 epoch 冻结/解冻对应层"""
        layers_to_unfreeze = sum(1 for s in self.schedule if s <= epoch)
        
        # 先冻结所有 Transformer 层
        for param in self.model.transformer.parameters():
            param.requires_grad = False
            
        # 解冻最后 N 层
        for i in range(min(layers_to_unfreeze, len(self.model.transformer.layers))):
            layer = self.model.transformer.layers[-(i+1)]
            for param in layer.parameters():
                param.requires_grad = True
                
        # 分类头始终可训练
        for param in self.model.classifier.parameters():
            param.requires_grad = True
            param.data.normal_(0, 0.01)
        for param in self.model.norm.parameters():
            param.requires_grad = True


def evaluate_cross_subject(model, dataset, subject_ids):
    """
    留一被试交叉验证（LOSO）
    
    Args:
        model: EEG 模型
        dataset: 完整数据集
        subject_ids: 所有被试 ID
    
    Returns:
        per_subject_acc: 每个被试的准确率
    """
    from sklearn.metrics import accuracy_score, f1_score
    
    results = {}
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model = model.to(device)
    
    for test_subj in np.unique(subject_ids):
        train_mask = subject_ids != test_subj
        test_mask = subject_ids == test_subj
        
        train_data = dataset.data[train_mask]
        train_labels = dataset.labels[train_mask]
        test_data = dataset.data[test_mask]
        test_labels = dataset.labels[test_mask]
        
        train_ds = torch.utils.data.TensorDataset(train_data, train_labels)
        test_ds = torch.utils.data.TensorDataset(test_data, test_labels)
        train_loader = DataLoader(train_ds, batch_size=64, shuffle=True)
        test_loader = DataLoader(test_ds, batch_size=64)
        
        # 重新初始化模型
        model.apply(lambda m: m.reset_parameters() if hasattr(m, 'reset_parameters') else None)
        optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4, weight_decay=0.01)
        criterion = nn.CrossEntropyLoss()
        
        unfreezer = ProgressiveUnfreezing(model, total_epochs=50)
        
        for epoch in range(50):
            model.train()
            unfreezer.apply_freeze(epoch)
            for batch_x, batch_y in train_loader:
                batch_x, batch_y = batch_x.to(device), batch_y.to(device)
                logits = model(batch_x)
                loss = criterion(logits, batch_y)
                optimizer.step()
                optimizer.zero_grad()
                loss.backward()
        
        # 测试
        model.eval()
        all_preds = []
        with torch.no_grad():
            for batch_x, _ in test_loader:
                batch_x = batch_x.to(device)
                logits = model(batch_x)
                all_preds.extend(logits.argmax(dim=1).cpu().numpy())
        
        acc = accuracy_score(test_labels.numpy(), all_preds)
        f1 = f1_score(test_labels.numpy(), all_preds, average='macro')
        results[f'subject_{test_subj}'] = {'accuracy': acc, 'f1': f1}
        
    return results


# ==================== 实际测试 ====================
if __name__ == "__main__":
    np.random.seed(42)
    torch.manual_seed(42)
    
    # 模拟驾驶员 EEG 数据
    # 19 通道 (标准 10-20 系统), 256 时间点 (2s @ 128Hz)
    num_subjects = 10
    samples_per_subject = 50
    num_channels = 19
    seq_len = 256
    
    # 生成模拟数据
    all_data = []
    all_labels = []
    all_subjects = []
    
    for subj_id in range(num_subjects):
        # 正常驾驶数据
        normal_data = np.random.randn(samples_per_subject // 2, num_channels, seq_len) * 0.5 + 0.3
        normal_labels = np.zeros(samples_per_subject // 2, dtype=int)
        
        # 认知分心数据 (增加 alpha 波段能量)
        distracted_data = np.random.randn(samples_per_subject // 2, num_channels, seq_len) * 0.5 + 0.3
        # 模拟 alpha 波段增强 (8-13Hz)
        t = np.arange(seq_len) / 128
        alpha_wave = np.sin(2 * np.pi * 10 * t)  # 10Hz alpha
        distracted_data += alpha_wave * 0.3
        
        all_data.extend([normal_data, distracted_data])
        all_labels.extend([normal_labels, np.ones(samples_per_subject // 2, dtype=int)])
        all_subjects.extend([subj_id] * samples_per_subject)
    
    eeg_data = np.concatenate(all_data, axis=0)
    labels = np.concatenate(all_labels, axis=0)
    subject_ids = np.array(all_subjects)
    
    print(f"数据集: {eeg_data.shape[0]} 样本, {num_channels} 通道, {seq_len} 时间点")
    print(f"被试数: {num_subjects}")
    
    # 创建数据集
    dataset = EEGDataset(eeg_data, labels, subject_ids)
    
    # 创建模型
    model = EEGTransformerEncoder(
        num_channels=num_channels,
        seq_len=seq_len,
        embed_dim=256,
        num_heads=8,
        num_layers=6,
        num_classes=2,
        patch_size=16
    )
    
    print(f"\n模型参数量: {sum(p.numel() for p in model.parameters()):,}")
    
    # LOSO 评估
    results = evaluate_cross_subject(model, dataset, subject_ids)
    
    print("\n===== 跨被试 LOSO 评估结果 =====")
    accs = []
    for subj, metrics in results.items():
        print(f"  {subj}: Acc={metrics['accuracy']:.4f}, F1={metrics['f1']:.4f}")
        accs.append(metrics['accuracy'])
    
    print(f"\n平均准确率: {np.mean(accs):.4f} ± {np.std(accs):.4f}")
    print(f"最佳被试: {max(accs):.4f}")
    print(f"最差被试: {min(accs):.4f}")