import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np

class NeurophysiologyGuidedGrouping:
    """
    神经生理引导的EEG通道分组
    
    基于脑科学先验将17通道分为3组：
    1. 前额认知组（Fp1, Fp2, F3, F4, Fz）→ 执行控制/注意力
    2. 中央运动组（C3, Cz, C4, T7, T8）→ 运动准备/体感
    3. 后部感知组（Pz, O1, O2, P3, P4, Oz, P7）→ 视觉处理/疲劳
    
    论文Section 2.1: Channel Grouping Strategy
    """
    
    # SEED-VIG 17通道国际10-20系统
    GROUPS = {
        'prefrontal': ['Fp1', 'Fp2', 'F3', 'F4', 'Fz'],  # 认知控制
        'central': ['C3', 'Cz', 'C4', 'T7', 'T8'],        # 运动/体感
        'posterior': ['Pz', 'O1', 'O2', 'P3', 'P4', 'Oz', 'P7']  # 视觉感知
    }
    
    def __init__(self):
        self.channel_indices = self._build_indices()
    
    def _build_indices(self):
        """构建通道索引映射"""
        # SEED-VIG通道顺序
        all_channels = ['Fp1', 'Fp2', 'F3', 'F4', 'Fz', 
                        'C3', 'Cz', 'C4', 'T7', 'T8',
                        'Pz', 'O1', 'O2', 'P3', 'P4', 'Oz', 'P7']
        
        indices = {}
        for group_name, channels in self.GROUPS.items():
            indices[group_name] = [all_channels.index(ch) for ch in channels 
                                   if ch in all_channels]
        return indices
    
    def group(self, eeg: torch.Tensor) -> dict:
        """
        Args:
            eeg: (B, 17, T) 17通道EEG信号
            
        Returns:
            dict: {
                'prefrontal': (B, 5, T),
                'central': (B, 5, T),
                'posterior': (B, 7, T)
            }
        """
        return {
            name: eeg[:, idx, :] 
            for name, idx in self.channel_indices.items()
        }


class SelectiveSSM(nn.Module):
    """
    选择性状态空间模型（Mamba架构）
    
    核心改进：输入自适应的选择机制
    - 传统SSM: 固定参数A, B, C, D
    - Mamba SSM: A(x), B(x), C(x), D(x) 随输入动态变化
    
    论文Section 2.2: Tri-Stream Selective State Space
    
    Args:
        d_model: 模型维度
        d_state: 状态空间维度
        d_conv: 卷积核大小
        expand: 扩展因子
    """
    
    def __init__(self, d_model: int = 64, d_state: int = 16,
                 d_conv: int = 4, expand: int = 2):
        super().__init__()
        self.d_model = d_model
        self.d_state = d_state
        self.d_conv = d_conv
        self.expand = expand
        self.d_inner = int(d_model * expand)
        
        # 输入投影
        self.in_proj = nn.Linear(d_model, self.d_inner * 2, bias=False)
        
        # 局部卷积
        self.conv1d = nn.Conv1d(
            self.d_inner, self.d_inner,
            kernel_size=d_conv, padding=d_conv - 1,
            groups=self.d_inner, bias=True
        )
        
        # x项目到Δ, B, C
        self.x_proj = nn.Linear(
            self.d_inner, d_state * 2 + self.d_inner, bias=False
        )
        
        # Δ, B, C → A, B, C
        self.dt_proj = nn.Linear(d_state, self.d_inner, bias=True)
        
        # A参数（log空间）
        self.A_log = nn.Parameter(
            torch.randn(self.d_inner, d_state) * 0.01
        )
        self.D = nn.Parameter(torch.ones(self.d_inner))
        
        # 输出投影
        self.out_proj = nn.Linear(self.d_inner, d_model, bias=False)
        
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Args:
            x: (B, L, d_model) 输入序列
            
        Returns:
            out: (B, L, d_model) 输出序列
        """
        B, L, _ = x.shape
        
        # 输入投影 + 分支
        x_proj = self.in_proj(x)  # (B, L, 2*d_inner)
        x1, x2 = x_proj.chunk(2, dim=-1)  # 各(B, L, d_inner)
        
        # 局部卷积 + 激活
        x1 = x1.transpose(1, 2)  # (B, d_inner, L)
        x1 = self.conv1d(x1)[:, :, :L]
        x1 = F.silu(x1)
        x1 = x1.transpose(1, 2)  # (B, L, d_inner)
        
        # 选择性参数生成
        x_ssm = self.x_proj(x1)  # (B, L, d_state*2 + d_inner)
        B_param = x_ssm[:, :, :self.d_state]  # (B, L, d_state)
        C_param = x_ssm[:, :, self.d_state:2*self.d_state]  # (B, L, d_state)
        delta = F.softplus(self.dt_proj(
            x_ssm[:, :, 2*self.d_state:]
        ))  # (B, L, d_inner)
        
        # 离散化A, B
        A = -torch.exp(self.A_log)  # (d_inner, d_state)
        A_discrete = torch.exp(delta.unsqueeze(-1) * A)  # (B, L, d_inner, d_state)
        B_discrete = delta.unsqueeze(-1) * B_param.unsqueeze(2)  # (B, L, d_inner, d_state)
        
        # SSM扫描（简化版，实际用parallel scan）
        h = torch.zeros(B, self.d_inner, self.d_state, device=x.device)
        outputs = []
        for t in range(L):
            h = A_discrete[:, t] * h + B_discrete[:, t] * x1[:, t].unsqueeze(-1)
            y = (h * C_param[:, t].unsqueeze(1)).sum(-1)  # (B, d_inner)
            outputs.append(y)
        
        y = torch.stack(outputs, dim=1)  # (B, L, d_inner)
        y = y + x1 * self.D
        
        # 门控 + 输出投影
        y = y * F.silu(x2)
        out = self.out_proj(y)
        
        return out


class TriStreamNTSM(nn.Module):
    """
    NTSM: 神经生理引导的三流选择性状态空间模型
    
    三流:
      Stream 1: EEG时域（通道分组 → 各组Mamba → 融合）
      Stream 2: EOG时域（眼动特征 → Mamba）
      Stream 3: 频域特征（δ/θ/α/β/γ → Mamba）
    
    论文Section 2: Full Architecture
    
    Args:
        eeg_channels: EEG通道数 (SEED-VIG: 17)
        eog_channels: EOG通道数 (2: 水平+垂直)
        d_model: 模型维度
        d_state: SSM状态维度
        num_classes: 分类数 (3: 清醒/轻度/重度)
    """
    
    def __init__(self, eeg_channels: int = 17, eog_channels: int = 2,
                 d_model: int = 64, d_state: int = 16, num_classes: int = 3):
        super().__init__()
        self.grouper = NeurophysiologyGuidedGrouping()
        
        # Stream 1: EEG时域（3个分组各一个SSM）
        self.eeg_embeds = nn.ModuleDict({
            name: nn.Linear(len(idx), d_model)
            for name, idx in self.grouper.channel_indices.items()
        })
        self.eeg_ssms = nn.ModuleDict({
            name: SelectiveSSM(d_model, d_state)
            for name in self.grouper.channel_indices
        })
        
        # Stream 2: EOG时域
        self.eog_embed = nn.Linear(eog_channels, d_model)
        self.eog_ssm = SelectiveSSM(d_model, d_state)
        
        # Stream 3: 频域特征（5个频段）
        self.freq_embed = nn.Linear(5 * eeg_channels, d_model)
        self.freq_ssm = SelectiveSSM(d_model, d_state)
        
        # 跨模态注意力融合
        self.cross_attn = nn.MultiheadAttention(
            d_model, num_heads=4, batch_first=True
        )
        
        # 分类头
        self.classifier = nn.Sequential(
            nn.LayerNorm(d_model),
            nn.Linear(d_model, 32),
            nn.GELU(),
            nn.Dropout(0.2),
            nn.Linear(32, num_classes)
        )
        
    def forward(self, eeg: torch.Tensor, eog: torch.Tensor,
                freq: torch.Tensor) -> torch.Tensor:
        """
        Args:
            eeg: (B, 17, L) EEG时域信号
            eog: (B, 2, L) EOG时域信号
            freq: (B, 17*5, L) 频域特征（5频段×17通道）
            
        Returns:
            logits: (B, num_classes)
        """
        # Stream 1: EEG时域（分组处理）
        groups = self.grouper.group(eeg)  # dict of (B, n, L)
        eeg_features = []
        for name, g in groups.items():
            g = g.transpose(1, 2)  # (B, L, n)
            g = self.eeg_embeds[name](g)  # (B, L, d_model)
            g = self.eeg_ssms[name](g)  # (B, L, d_model)
            eeg_features.append(g.mean(dim=1))  # (B, d_model)
        eeg_feat = torch.stack(eeg_features, dim=1)  # (B, 3, d_model)
        
        # Stream 2: EOG时域
        eog_in = eog.transpose(1, 2)  # (B, L, 2)
        eog_in = self.eog_embed(eog_in)  # (B, L, d_model)
        eog_out = self.eog_ssm(eog_in)  # (B, L, d_model)
        eog_feat = eog_out.mean(dim=1).unsqueeze(1)  # (B, 1, d_model)
        
        # Stream 3: 频域
        freq_in = freq.transpose(1, 2)  # (B, L, 85)
        freq_in = self.freq_embed(freq_in)  # (B, L, d_model)
        freq_out = self.freq_ssm(freq_in)  # (B, L, d_model)
        freq_feat = freq_out.mean(dim=1).unsqueeze(1)  # (B, 1, d_model)
        
        # 跨模态注意力融合
        all_features = torch.cat([eeg_feat, eog_feat, freq_feat], dim=1)  # (B, 5, d_model)
        attn_out, _ = self.cross_attn(all_features, all_features, all_features)
        fused = attn_out.mean(dim=1)  # (B, d_model)
        
        # 分类
        return self.classifier(fused)


# ============ 完整测试代码 ============
if __name__ == "__main__":
    import time
    from sklearn.metrics import accuracy_score
    
    # SEED-VIG参数
    EEG_CHANNELS = 17
    EOG_CHANNELS = 2
    SEQ_LEN = 1200  # 5分钟 @ 4Hz
    NUM_CLASSES = 3  # 清醒/轻度疲劳/重度疲劳
    
    # 初始化模型
    model = TriStreamNTSM(
        eeg_channels=EEG_CHANNELS,
        eog_channels=EOG_CHANNELS,
        d_model=64,
        d_state=16,
        num_classes=NUM_CLASSES
    )
    
    total_params = sum(p.numel() for p in model.parameters())
    print(f"NTSM参数量: {total_params:,}")
    
    # 模拟EEG/EOG数据
    np.random.seed(42)
    
    def generate_eeg_eog(state: str, batch=1):
        """根据疲劳状态生成模拟EEG+EOG"""
        t = np.linspace(0, 300, SEQ_LEN)
        
        if state == 'alert':
            # 清醒：高频低幅β波 + 频繁眼动
            eeg = np.random.randn(batch, EEG_CHANNELS, SEQ_LEN) * 5
            eeg += 3 * np.sin(2*np.pi*25*t + np.random.rand())  # β波
            eog = np.random.randn(batch, EOG_CHANNELS, SEQ_LEN) * 20
            eog += 50 * (np.random.rand(batch, EOG_CHANNELS, SEQ_LEN) > 0.95)  # 扫视
            
        elif state == 'mild':
            # 轻度疲劳：α波增强 + 眼动减少
            eeg = np.random.randn(batch, EEG_CHANNELS, SEQ_LEN) * 5
            eeg += 8 * np.sin(2*np.pi*10*t)  # α波增多
            eog = np.random.randn(batch, EOG_CHANNELS, SEQ_LEN) * 10
            eog += 30 * (np.random.rand(batch, EOG_CHANNELS, SEQ_LEN) > 0.98)
            
        else:  # severe
            # 重度疲劳：θ/δ波主导 + 慢眼动
            eeg = np.random.randn(batch, EEG_CHANNELS, SEQ_LEN) * 3
            eeg += 15 * np.sin(2*np.pi*4*t)  # θ波
            eeg += 10 * np.sin(2*np.pi*2*t)  # δ波
            eog = np.random.randn(batch, EOG_CHANNELS, SEQ_LEN) * 5
            eog += 20 * np.sin(2*np.pi*0.5*t)  # 慢眼动
        
        # 频域特征：5频段×17通道
        freq = np.random.rand(batch, EEG_CHANNELS * 5, SEQ_LEN).astype(np.float32)
        if state == 'alert':
            freq[:, ::5, :] *= 2.0  # β频段高
        elif state == 'mild':
            freq[:, 1::5, :] *= 2.0  # α频段高
        else:
            freq[:, 3::5, :] *= 2.0  # δ频段高
        
        return (torch.tensor(eeg, dtype=torch.float32),
                torch.tensor(eog, dtype=torch.float32),
                torch.tensor(freq, dtype=torch.float32))
    
    # 生成数据
    print("\n--- 生成模拟数据 ---")
    states = ['alert', 'mild', 'severe']
    all_eeg, all_eog, all_freq, all_labels = [], [], [], []
    
    for i, state in enumerate(states):
        eeg, eog, freq = generate_eeg_eog(state, batch=10)
        all_eeg.append(eeg)
        all_eog.append(eog)
        all_freq.append(freq)
        all_labels.extend([i] * 10)
    
    eeg_batch = torch.cat(all_eeg)
    eog_batch = torch.cat(all_eog)
    freq_batch = torch.cat(all_freq)
    labels = torch.tensor(all_labels)
    
    print(f"EEG: {eeg_batch.shape}")
    print(f"EOG: {eog_batch.shape}")
    print(f"Freq: {freq_batch.shape}")
    
    # 训练流程
    print("\n--- 训练流程 ---")
    criterion = nn.CrossEntropyLoss()
    optimizer = torch.optim.AdamW(model.parameters(), lr=2e-4, weight_decay=0.01)
    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=50)
    
    model.train()
    for epoch in range(5):
        optimizer.zero_grad()
        logits = model(eeg_batch, eog_batch, freq_batch)
        loss = criterion(logits, labels)
        loss.backward()
        optimizer.step()
        scheduler.step()
        
        acc = accuracy_score(labels.numpy(), 
                            torch.argmax(logits, dim=1).detach().numpy())
        print(f"Epoch {epoch+1}: Loss={loss.item():.4f}, Acc={acc:.2%}, LR={scheduler.get_last_lr()[0]:.6f}")
    
    # 评估
    print("\n--- 评估 ---")
    model.eval()
    with torch.no_grad():
        logits = model(eeg_batch, eog_batch, freq_batch)
        preds = torch.argmax(logits, dim=1)
        
        for i, state in enumerate(states):
            mask = (labels == i)
            acc = (preds[mask] == i).float().mean()
            print(f"  {state:8s}: Acc={acc:.2%} ({mask.sum().item()} samples)")
    
    # 推理速度
    with torch.no_grad():
        start = time.time()
        for _ in range(20):
            _ = model(eeg_batch[:1], eog_batch[:1], freq_batch[:1])
        elapsed = (time.time() - start) / 20
        print(f"\n单样本推理延迟: {elapsed*1000:.1f}ms")
        print(f"理论帧率: {1/elapsed:.0f} FPS")
        print(f"复杂度: O(n) 线性序列建模")