# 示意代码，不可直接运行（论文未开源，仅为融合/元学习/联邦思想的示意实现；论文实际采用 CNN–BiLSTM 混合编码器 + Transformer 融合 + DTW 时序对齐，与本代码结构不同）
import torch
import torch.nn as nn

class CrossModalTransformerFusion(nn.Module):
    """
    跨模态Transformer融合层
    
    论文核心：捕获模态间细粒度时空依赖
    如：EEG theta波变化 ↔ 眼动注视丢失
    """
    
    def __init__(self, d_model: int = 128, 
                 n_heads: int = 8,
                 n_modalities: int = 5):
        super().__init__()
        # 各模态投影到统一维度
        self.modality_proj = nn.ModuleList([
            nn.Linear(d_model, d_model) 
            for _ in range(n_modalities)
        ])
        
        # 跨模态注意力
        self.cross_attention = nn.MultiheadAttention(
            embed_dim=d_model,
            num_heads=n_heads,
            dropout=0.1,
            batch_first=True
        )
        
        # 前馈网络
        self.ffn = nn.Sequential(
            nn.Linear(d_model, d_model * 4),
            nn.GELU(),
            nn.Dropout(0.1),
            nn.Linear(d_model * 4, d_model),
        )
        
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        
    def forward(self, modality_features: list) -> torch.Tensor:
        """
        Args:
            modality_features: [B, T, D] × n_modalities
            
        Returns:
            fused: [B, T, D] 融合后特征
        """
        # 投影各模态
        projected = [
            proj(feat) for proj, feat in zip(
                self.modality_proj, modality_features
            )
        ]
        
        # 拼接所有模态作为序列
        # [n_modalities × B, T, D]
        all_features = torch.stack(projected, dim=1)  # [B, M, T, D]
        B, M, T, D = all_features.shape
        all_features = all_features.view(B, M * T, D)
        
        # 自注意力（跨模态+跨时间）
        attended, _ = self.cross_attention(
            all_features, all_features, all_features
        )
        
        # 残差+归一化
        x = self.norm1(all_features + attended)
        
        # FFN
        ffn_out = self.ffn(x)
        x = self.norm2(x + ffn_out)
        
        # 分离各模态并融合
        x = x.view(B, M, T, D)
        fused = x.mean(dim=1)  # 模态平均
        
        return fused


class AdaptiveMetaLearner(nn.Module):
    """
    个性化元学习：MAML-based
    
    论文Section 3.3:
    - ≤5个样本适配新驾驶员
    - 元学习初始化参数
    """
    
    def __init__(self, model: nn.Module, 
                 inner_lr: float = 0.01,
                 outer_lr: float = 0.001,
                 inner_steps: int = 5):
        super().__init__()
        self.model = model
        self.inner_lr = inner_lr
        self.outer_lr = outer_lr
        self.inner_steps = inner_steps
        
    def meta_train_step(self, support_set, query_set):
        """
        MAML元训练步骤
        
        Args:
            support_set: 新驾驶员的少量样本（≤5个）
            query_set: 查询集用于评估
        """
        # 1. 内循环：在support set上快速适配
        fast_weights = {}
        for name, param in self.model.named_parameters():
            fast_weights[name] = param.clone()
        
        for step in range(self.inner_steps):
            support_loss = self._compute_loss(
                self.model, support_set, fast_weights
            )
            grads = torch.autograd.grad(
                support_loss, fast_weights.values()
            )
            fast_weights = {
                name: w - self.inner_lr * g
                for (name, w), g in zip(fast_weights.items(), grads)
            }
        
        # 2. 外循环：在query set上计算meta-loss
        query_loss = self._compute_loss(
            self.model, query_set, fast_weights
        )
        
        return query_loss
    
    def _compute_loss(self, model, data, weights):
        """用fast weights计算损失"""
        x, y = data
        # 前向传播用fast weights
        output = model.forward_with_weights(x, weights)
        return nn.functional.cross_entropy(output, y)


class FederatedOptimizer:
    """
    联邦优化器
    
    论文Section 3.4:
    - 去中心化训练
    - 自适应梯度压缩
    - 非IID数据处理
    """
    
    def __init__(self, global_model: nn.Module,
                 n_clients: int,
                 compression_ratio: float = 0.1):
        self.global_model = global_model
        self.n_clients = n_clients
        self.compression_ratio = compression_ratio
        
    def client_update(self, client_id: int, 
                      local_data, epochs: int = 5) -> dict:
        """客户端本地训练"""
        local_model = copy.deepcopy(self.global_model)
        local_model.train()
        
        optimizer = torch.optim.Adam(
            local_model.parameters(), lr=1e-3
        )
        
        for epoch in range(epochs):
            for batch in local_data:
                optimizer.zero_grad()
                loss = self._compute_loss(local_model, batch)
                loss.backward()
                optimizer.step()
        
        # 梯度压缩（仅传top-k%参数差）
        global_state = self.global_model.state_dict()
        local_state = local_model.state_dict()
        
        updates = {
            k: local_state[k] - global_state[k]
            for k in global_state
        }
        
        # Top-k稀疏化
        compressed = self._compress(updates)
        return compressed
    
    def _compress(self, updates: dict) -> dict:
        """Top-k梯度压缩"""
        all_vals = torch.cat([
            v.flatten() for v in updates.values()
        ])
        k = int(len(all_vals) * self.compression_ratio)
        topk_vals, topk_idx = torch.topk(all_vals.abs(), k)
        
        mask = torch.zeros_like(all_vals)
        mask[topk_idx] = 1.0
        
        compressed = {}
        offset = 0
        for name, update in updates.items():
            n = update.numel()
            compressed[name] = update * mask[offset:offset+n].view_as(update)
            offset += n
        
        return compressed
    
    def aggregate(self, client_updates: list, 
                  client_weights: list) -> dict:
        """FedAvg聚合"""
        global_state = self.global_model.state_dict()
        aggregated = {k: torch.zeros_like(v) 
                       for k, v in global_state.items()}
        
        for update, weight in zip(client_updates, client_weights):
            for k in aggregated:
                aggregated[k] += update[k] * weight
        
        for k in global_state:
            global_state[k] += aggregated[k]
        
        self.global_model.load_state_dict(global_state)


if __name__ == "__main__":
    import copy
    
    # 测试跨模态融合
    fusion = CrossModalTransformerFusion(d_model=128, n_modalities=5)
    
    # 模拟5个模态特征
    modalities = [torch.randn(4, 30, 128) for _ in range(5)]
    
    fused = fusion(modalities)
    print(f"输入: 5个模态 × {modalities[0].shape}")
    print(f"融合输出: {fused.shape}")
    print(f"参数量: {sum(p.numel() for p in fusion.parameters()):,}")