#!/usr/bin/env python3
"""
联邦学习模块 — 虫群节点间权重聚合

三档模型(lite/standard/pro) + FedAvg跨档位聚合
支持: 权重同步、增量diff、弹性加入/退出
"""
import numpy as np
import base64
import io
import time
import copy
from typing import Dict, List, Optional, Tuple
from pathlib import Path


# ============================================================
# 三档模型配置
# ============================================================
TIER_CONFIG = {
    'lite': {'embed_dim': 128, 'hidden': 512, 'layers': 2, 'max_vocab': 10000},
    'standard': {'embed_dim': 256, 'hidden': 1024, 'layers': 2, 'max_vocab': 10000},
    'pro': {'embed_dim': 256, 'hidden': 2048, 'layers': 3, 'max_vocab': 20000},
}

AREA_NAMES = {
    'sensory': '感觉区', 'memory': '记忆区', 'association': '联想区',
    'motor': '运动区', 'prefrontal': '前额叶', 'thalamus': '丘脑',
}


class WeightCodec:
    """权重编解码器 — numpy ↔ base64字符串"""
    
    @staticmethod
    def encode(weights: Dict[str, np.ndarray]) -> str:
        """权重字典 → base64字符串"""
        buf = io.BytesIO()
        np.savez_compressed(buf, **weights)
        return base64.b64encode(buf.getvalue()).decode('ascii')
    
    @staticmethod
    def decode(data: str) -> Dict[str, np.ndarray]:
        """base64字符串 → 权重字典"""
        raw = base64.b64decode(data)
        buf = io.BytesIO(raw)
        return dict(np.load(buf))


class FedAvgAggregator:
    """FedAvg聚合器 — 多节点权重平均"""
    
    def __init__(self):
        self._weight_buffer: Dict[str, List[Dict[str, np.ndarray]]] = {}
        self._node_versions: Dict[str, int] = {}
    
    def submit(self, node_id: str, area: str, 
               weights: Dict[str, np.ndarray], version: int):
        """提交节点权重"""
        key = f"{node_id}:{area}"
        if area not in self._weight_buffer:
            self._weight_buffer[area] = []
        self._weight_buffer[area].append(weights)
        self._node_versions[key] = version
    
    def aggregate(self, area: str) -> Optional[Dict[str, np.ndarray]]:
        """FedAvg聚合: W_avg = ΣW_i / N"""
        if area not in self._weight_buffer or not self._weight_buffer[area]:
            return None
        
        weight_list = self._weight_buffer[area]
        n = len(weight_list)
        
        # 取第一个作为模板
        result = {}
        for key in weight_list[0]:
            tensors = [w[key] for w in weight_list if key in w]
            if tensors:
                # 零填充升维(不同档位)
                max_shape = max(t.shape for t in tensors)
                padded = []
                for t in tensors:
                    if t.shape == max_shape:
                        padded.append(t)
                    else:
                        p = np.zeros(max_shape, dtype=np.float32)
                        slices = tuple(slice(0, s) for s in t.shape)
                        p[slices] = t
                        padded.append(p)
                result[key] = np.mean(padded, axis=0)
        
        # 清空缓冲
        self._weight_buffer[area] = []
        return result
    
    def get_version(self, node_id: str, area: str) -> int:
        return self._node_versions.get(f"{node_id}:{area}", 0)


class DiffSync:
    """增量同步 — 只传权重变化量"""
    
    @staticmethod
    def compute_diff(old: Dict[str, np.ndarray], 
                     new: Dict[str, np.ndarray],
                     threshold: float = 1e-6) -> Dict[str, np.ndarray]:
        """计算权重差异"""
        diff = {}
        for key in new:
            if key in old:
                d = new[key] - old[key]
                if np.max(np.abs(d)) > threshold:
                    diff[key] = d
            else:
                diff[key] = new[key]
        return diff
    
    @staticmethod
    def apply_diff(base: Dict[str, np.ndarray],
                   diff: Dict[str, np.ndarray]) -> Dict[str, np.ndarray]:
        """应用差异到基础权重"""
        result = {k: v.copy() for k, v in base.items()}
        for key, delta in diff.items():
            if key in result:
                result[key] = result[key] + delta
            else:
                result[key] = delta
        return result


class FederationNode:
    """联邦学习节点 — 管理本地训练和远程同步"""
    
    def __init__(self, node_id: str, tier: str = 'lite',
                 areas: Optional[List[str]] = None):
        self.node_id = node_id
        self.tier = tier
        self.areas = areas or list(AREA_NAMES.keys())
        
        # 本地权重版本
        self._versions: Dict[str, int] = {a: 0 for a in self.areas}
        
        # 上次同步的权重快照(用于diff计算)
        self._snapshots: Dict[str, Dict[str, np.ndarray]] = {}
        
        # 统计
        self.stats = {
            'local_updates': 0,
            'sync_sent': 0,
            'sync_received': 0,
            'last_sync': None,
        }
    
    def after_train(self, area: str, weights: Dict[str, np.ndarray]):
        """本地训练后更新版本"""
        self._versions[area] += 1
        self.stats['local_updates'] += 1
    
    def get_sync_payload(self, area: str, 
                         weights: Dict[str, np.ndarray],
                         use_diff: bool = True) -> Dict:
        """准备同步数据"""
        version = self._versions[area]
        
        if use_diff and area in self._snapshots:
            diff = DiffSync.compute_diff(self._snapshots[area], weights)
            payload_weights = diff
            mode = 'diff'
        else:
            payload_weights = weights
            mode = 'full'
        
        # 保存快照
        self._snapshots[area] = {k: v.copy() for k, v in weights.items()}
        
        self.stats['sync_sent'] += 1
        self.stats['last_sync'] = time.time()
        
        return {
            'node_id': self.node_id,
            'area': area,
            'version': version,
            'tier': self.tier,
            'mode': mode,
            'weights': WeightCodec.encode(payload_weights),
        }
    
    def apply_received(self, area: str, 
                       weights: Dict[str, np.ndarray],
                       mode: str = 'full',
                       base_weights: Optional[Dict[str, np.ndarray]] = None):
        """应用接收到的聚合权重"""
        if mode == 'diff' and base_weights:
            result = DiffSync.apply_diff(base_weights, weights)
        else:
            result = weights
        
        self.stats['sync_received'] += 1
        return result