"""
EF-SCConv-CapsGCN: EEG-fNIRS Spatial-Channel Convolution + Capsule GCN

论文核心模型复现
参考：Expert Systems with Applications, 2026
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from typing import Tuple

class SCConv(nn.Module):
    """
    Spatial-Channel Convolution (SCConv)
    
    同时提取空间维（脑区位置）和通道维（特征类型）的特征
    - 空间维：脑电极位置的空间关系
    - 通道维：不同频段的功率特征
    """
    
    def __init__(self, in_channels: int, out_channels: int,
                 n_eeg_channels: int = 32, n_fnirs_channels: int = 16):
        super().__init__()
        
        # 空间卷积（沿通道方向）
        self.spatial_conv = nn.Conv2d(
            in_channels, out_channels,
            kernel_size=(n_eeg_channels + n_fnirs_channels, 1),
            stride=1, padding=(0, 0)
        )
        
        # 通道卷积（沿特征方向）
        self.channel_conv = nn.Conv2d(
            out_channels, out_channels,
            kernel_size=(1, 3), stride=1, padding=(0, 1),
            groups=out_channels  # 深度可分离
        )
        
        self.bn = nn.BatchNorm2d(out_channels)
        self.act = nn.GELU()
        
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Args:
            x: (B, C, T) where C = EEG_ch + fNIRS_ch, T = time_steps
        Returns:
            out: (B, out_C, 1, T)
        """
        x = x.unsqueeze(2)  # (B, C, 1, T)
        x = self.spatial_conv(x)  # (B, out_C, 1, T)
        x = self.channel_conv(x)  # (B, out_C, 1, T)
        x = self.bn(x)
        x = self.act(x)
        return x.squeeze(2)  # (B, out_C, T)


class CapsuleGCN(nn.Module):
    """
    Capsule Graph Convolutional Network
    
    将脑区视为图节点，用图卷积建模脑区间功能连接
    胶囊网络替代传统 FC 层，保留空间层次关系
    """
    
    def __init__(self, n_features: int, n_regions: int = 7,
                 n_classes: int = 3, capsule_dim: int = 16):
        super().__init__()
        
        self.n_regions = n_regions
        self.capsule_dim = capsule_dim
        
        # 脑区邻接矩阵（基于先验解剖知识）
        self.register_buffer(
            'adj_matrix',
            self._create_brain_region_adjacency(n_regions)
        )
        
        # 图卷积层
        self.gc1 = nn.Linear(n_features, 64)
        self.gc2 = nn.Linear(64, capsule_dim)
        
        # 胶囊路由
        self.routing_iterations = 3
        self.W_route = nn.Parameter(
            torch.randn(n_regions, capsule_dim, capsule_dim) * 0.01
        )
        
        # 分类器
        self.classifier = nn.Linear(n_regions * capsule_dim, n_classes)
        
    def _create_brain_region_adjacency(self, n: int) -> torch.Tensor:
        """创建脑区邻接矩阵"""
        adj = torch.zeros(n, n)
        regions = ['frontal', 'central', 'parietal', 'occipital',
                   'temporal_l', 'temporal_r', 'motor']
        # 前额叶连接最广
        adj[0, [1, 5, 6, 4]] = 1
        adj[1, [0, 2, 6]] = 1
        adj[2, [1, 3]] = 1
        adj[3, [2]] = 1
        adj[4, [0]] = 1
        adj[5, [0]] = 1
        adj[6, [0, 1]] = 1
        adj = (adj + adj.T).clamp(max=1)
        adj.fill_diagonal_(1)
        # 对称归一化
        degree = adj.sum(dim=1, keepdim=True)
        adj_norm = adj / (degree + 1e-8)
        return adj_norm
        
    def graph_conv(self, x: torch.Tensor) -> torch.Tensor:
        """图卷积"""
        x = torch.matmul(self.adj_matrix, x)  # 邻域聚合
        return x
        
    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        Args:
            x: (B, n_regions, n_features)
        Returns:
            logits: (B, n_classes)
            capsules: (B, n_regions, capsule_dim)
        """
        B, N, F_in = x.shape
        
        # GCN Layer 1
        h = self.graph_conv(x)
        h = F.gelu(self.gc1(h))
        
        # GCN Layer 2
        h = self.graph_conv(h)
        h = F.gelu(self.gc2(h))  # (B, N, capsule_dim)
        
        # 动态路由胶囊
        capsules = self._dynamic_routing(h)
        
        # 展平并分类
        flat = capsules.view(B, -1)
        logits = self.classifier(flat)
        
        return logits, capsules
    
    def _dynamic_routing(self, h: torch.Tensor) -> torch.Tensor:
        """动态路由算法（简化版）"""
        B, N, D = h.shape
        
        # 初始化 logits
        b = torch.zeros(B, N, N, device=h.device)
        
        for i in range(self.routing_iterations):
            # softmax over j
            c = F.softmax(b, dim=2)  # (B, N, N)
            
            # 计算预测向量
            s = torch.einsum('bij,jkl->bikl', c, self.W_route)
            s = s.sum(dim=1)  # (B, N, D)
            
            # squash
            v = self._squash(s)
            
            # 更新 logits
            if i < self.routing_iterations - 1:
                b = b + torch.einsum('bik,bjk->bij', v, h)
        
        return v
    
    @staticmethod
    def _squash(s: torch.Tensor) -> torch.Tensor:
        """Squash 激活函数"""
        sq = torch.norm(s, dim=-1, keepdim=True)
        return s * sq / (1 + sq**2)


class EFSCConvCapsGCN(nn.Module):
    """
    完整的 EF-SCConv-CapsGCN 模型
    
    输入：EEG (B, 32, T) + fNIRS (B, 16, T)
    输出：警觉等级 (B, 3) [清醒, 轻度疲劳, 重度疲劳]
    """
    
    def __init__(self, n_eeg: int = 32, n_fnirs: int = 16,
                 n_classes: int = 3, n_brain_regions: int = 7):
        super().__init__()
        
        # EEG 分支
        self.eeg_scconv = SCConv(1, 32, n_eeg_channels=n_eeg)
        self.eeg_pool = nn.AdaptiveAvgPool1d(1)
        
        # fNIRS 分支
        self.fnirs_scconv = SCConv(1, 16, n_eeg_channels=n_fnirs)
        self.fnirs_pool = nn.AdaptiveAvgPool1d(1)
        
        # 融合
        self.fusion = nn.Linear(32 + 16, n_brain_regions * 32)
        
        # CapsGCN
        self.caps_gcn = CapsuleGCN(
            n_features=32, n_regions=n_brain_regions,
            n_classes=n_classes, capsule_dim=16
        )
        
    def forward(self, eeg: torch.Tensor, fnirs: torch.Tensor):
        """
        Args:
            eeg: (B, 32, T)
            fnirs: (B, 16, T)
        """
        # 分支提取
        eeg_feat = self.eeg_scconv(eeg.unsqueeze(1))  # (B, 32, T)
        eeg_feat = self.eeg_pool(eeg_feat).squeeze(-1)  # (B, 32)
        
        fnirs_feat = self.fnirs_scconv(fnirs.unsqueeze(1))  # (B, 16, T)
        fnirs_feat = self.fnirs_pool(fnirs_feat).squeeze(-1)  # (B, 16)
        
        # 融合
        fused = torch.cat([eeg_feat, fnirs_feat], dim=1)  # (B, 48)
        region_feat = self.fusion(fused)  # (B, N*32)
        region_feat = region_feat.view(-1, 7, 32)  # (B, 7, 32)
        
        # CapsGCN 分类
        logits, capsules = self.caps_gcn(region_feat)
        
        return logits, capsules


# 测试
if __name__ == "__main__":
    model = EFSCConvCapsGCN(n_eeg=32, n_fnirs=16, n_classes=3)
    
    # 模拟输入
    eeg = torch.randn(4, 32, 100)   # batch=4, 32通道, 100时间步
    fnirs = torch.randn(4, 16, 100)  # batch=4, 16通道, 100时间步
    
    logits, capsules = model(eeg, fnirs)
    print(f"Logits shape: {logits.shape}")
    print(f"Capsules shape: {capsules.shape}")
    print(f"Prediction: {torch.argmax(logits, dim=-1)}")
    print(f"Total params: {sum(p.numel() for p in model.parameters()):,}")

"""
认知分心多模态融合：EEG 指标 + 眼动指标 + 行为指标
"""

from dataclasses import dataclass
from typing import Optional
import numpy as np

@dataclass
class CognitiveDistractionInput:
    """认知分心检测输入"""
    # EEG 指标
    eeg_theta_alpha_ratio: Optional[float]  # θ/α 比率（前额叶）
    eeg_frontal_theta: Optional[float]     # 前额叶 θ 功率
    # fNIRS 指标
    fnirs_hbo_slope: Optional[float]       # 前额叶 HbO 变化斜率
    # 眼动指标
    gaze_entropy: Optional[float]           # 视线熵（越高越分散）
    blink_rate: Optional[float]             # 眨眼频率
    saccade_frequency: Optional[float]     # 扫视频率
    pupil_diameter_var: Optional[float]    # 瞳孔直径方差
    # 行为指标
    steering_entropy: Optional[float]      # 转向熵
    lane_deviation: Optional[float]        # 车道偏离
    speed_variation: Optional[float]       # 速度变异系数

class CognitiveDistractionFusion:
    """多模态认知分心融合判定"""
    
    def __init__(self):
        self.weights = {
            'eeg': 0.25,      # EEG 权重
            'fnirs': 0.15,    # fNIRS 权重（座舱部署更难）
            'eye': 0.40,      # 眼动权重（最实用）
            'behavior': 0.20, # 行为权重
        }
        self.thresholds = {
            'normal': 0.3,
            'mild': 0.6,
            'severe': 0.8,
        }
    
    def assess(self, inputs: CognitiveDistractionInput) -> dict:
        """评估认知分心等级"""
        scores = []
        
        # EEG 评分
        if inputs.eeg_theta_alpha_ratio is not None:
            # θ/α > 1.5 提示认知负荷过高
            eeg_score = min(1.0, inputs.eeg_theta_alpha_ratio / 3.0)
            scores.append(('eeg', eeg_score))
        
        # fNIRS 评分
        if inputs.fnirs_hbo_slope is not None:
            # HbO 斜率 < 0 提示认知资源耗竭
            fnirs_score = max(0, min(1.0, 0.5 - inputs.fnirs_hbo_slope * 0.1))
            scores.append(('fnirs', fnirs_score))
        
        # 眼动评分
        if inputs.gaze_entropy is not None:
            # 视线熵 > 0.7 提示注意力分散
            eye_score = min(1.0, inputs.gaze_entropy / 1.0)
            scores.append(('eye', eye_score))
        
        # 行为评分
        if inputs.steering_entropy is not None:
            beh_score = min(1.0, inputs.steering_entropy / 0.5)
            scores.append(('behavior', beh_score))
        
        # 加权融合
        total_score = 0
        total_weight = 0
        for modality, score in scores:
            w = self.weights.get(modality, 0.1)
            total_score += w * score
            total_weight += w
        
        if total_weight > 0:
            final_score = total_score / total_weight
        else:
            final_score = 0
        
        # 等级判定
        if final_score >= self.thresholds['severe']:
            level = 'SEVERE'
        elif final_score >= self.thresholds['mild']:
            level = 'MILD'
        elif final_score >= self.thresholds['normal']:
            level = 'SUSPECTED'
        else:
            level = 'NORMAL'
        
        return {
            'distraction_score': final_score,
            'level': level,
            'modality_scores': dict(scores),
        }


# 测试
if __name__ == "__main__":
    fusion = CognitiveDistractionFusion()
    
    # 场景1：正常驾驶
    normal_input = CognitiveDistractionInput(
        eeg_theta_alpha_ratio=0.8,
        fnirs_hbo_slope=0.02,
        gaze_entropy=0.3,
        blink_rate=15,
        saccade_frequency=2.0,
        steering_entropy=0.1,
    )
    result = fusion.assess(normal_input)
    print(f"Normal: {result['level']}, score={result['distraction_score']:.2f}")
    
    # 场景2：认知分心（走神思考问题）
    distracted_input = CognitiveDistractionInput(
        eeg_theta_alpha_ratio=2.5,
        fnirs_hbo_slope=-0.05,
        gaze_entropy=0.85,
        blink_rate=25,
        saccade_frequency=0.5,
        steering_entropy=0.35,
    )
    result2 = fusion.assess(distracted_input)
    print(f"Distracted: {result2['level']}, score={result2['distraction_score']:.2f}")
    print(f"  EEG: {result2['modality_scores'].get('eeg', 'N/A')}")
    print(f"  Eye: {result2['modality_scores'].get('eye', 'N/A')}")