src/core/brain.py

#!/usr/bin/env python3
"""
虫群Brain v2.0 — AreaNetwork适配层

包装Meta Model的AreaNetwork(6区×3柱=18柱)，补充chat/stats/get_area_weights等接口
让node_server和chat_engine无需修改即可使用完整版Meta Model引擎
"""
import numpy as np
import base64
import io
import os
from typing import Dict, Optional

from .functional_area import AreaNetwork, AreaType
from .language_decoder import LanguageDecoder


class Brain:
    """
    虫群Brain v2.0 — AreaNetwork适配器
    
    支持两种推理模式:
      - 原始模式: 全链路forward(6区随机权重, 信号衰减严重)
      - 映射模式: W_s(sensory直接映射) + P(motor投影), 跳过随机中间层
    
    映射模式在训练后自动启用
    """
    
    def __init__(self, dim: int = 75, config: dict = None):
        self.dim = dim
        self._config = config  # 保存配置供后续使用
        # AreaNetwork需要 {'sensory': {...}, ...} 格式
        area_config = config.get('areas', config) if config else None
        self._net = AreaNetwork(config=area_config)
        self._forward_count = 0
        # 默认冻结: forward只做推理，learn/train时临时解冻
        self._net.freeze()
        # 训练映射矩阵(加载后启用映射模式)
        self._W_s = None  # (75, 300) sensory直接映射
        self._P = None    # (300, 75) motor投影
        self._load_mappings()
        # QA检索索引(训练数据的问题向量→答案)
        self._qa_vecs = None   # (N, 75) 问题编码向量
        self._qa_answers = []  # 对应答案列表
        self._load_qa_index()
        # 语言解码层 — motor输出→自回归token生成
        self._decoder = None
        self._init_decoder()
    
    def _init_decoder(self):
        """初始化语言解码层，加载词表和已训练权重"""
        try:
            vocab_path = os.path.join(os.path.dirname(__file__), '..', '..', 'models', 'vocab75_clean_v2.pkl')
            # 尝试json格式(优先)
            json_path = vocab_path.replace('.pkl', '.json')
            chars = None
            if os.path.exists(json_path):
                import json as _json
                with open(json_path, 'r', encoding='utf-8') as f:
                    vocab_data = _json.load(f)
                chars = list(vocab_data) if isinstance(vocab_data, list) else list(vocab_data.keys())
            elif os.path.exists(vocab_path):
                import pickle
                with open(vocab_path, 'rb') as f:
                    vocab_data = pickle.load(f)
                chars = list(vocab_data.keys()) if isinstance(vocab_data, dict) else list(vocab_data)
            else:
                # 尝试clean词表
                alt_path = vocab_path.replace('vocab75_clean_v2', 'vocab75_clean')
                alt_json = alt_path.replace('.pkl', '_words.json')
                if os.path.exists(alt_json):
                    import json as _json
                    with open(alt_json, 'r', encoding='utf-8') as f:
                        chars = _json.load(f)
                elif os.path.exists(alt_path.replace('.pkl', '.npz')):
                    # 从npz加载但需要词表
                    pass
            if chars is None:
                print('[Brain] 未找到词表，跳过解码层初始化')
                self._decoder = None
                return
            motor_dim = self.config.get('motor_output_dim', 300)
            hidden_dim = self.config.get('areas', {}).get('motor', {}).get('hidden_dim', 300)
            self._decoder = LanguageDecoder(motor_dim=motor_dim, hidden_dim=hidden_dim, vocab_size=len(chars) + 3)
            self._decoder.set_vocab(chars)
            # 尝试加载已训练权重
            dec_path = os.path.join(os.path.dirname(__file__), '..', '..', 'models', 'decoder_weights.npz')
            if os.path.exists(dec_path):
                self._decoder.load(dec_path)
                print(f"[Brain] 语言解码层已加载训练权重")
            else:
                print(f"[Brain] 语言解码层已初始化(vocab={len(chars)}字, 未训练)")
        except Exception as e:
            print(f"[Brain] 语言解码层初始化失败: {e}")
            self._decoder = None
        
    @property
    def areas(self) -> Dict:
        """兼容旧接口: 返回 {区名: FunctionalArea}"""
        return {at.value: area for at, area in self._net.areas.items()}
    
    def _load_mappings(self):
        """加载训练映射矩阵(W_s和P), 存在则启用映射模式"""
        try:
            import os
            path = os.path.join(os.path.dirname(__file__), '..', '..', 'models', 'trained_mappings.npz')
            if os.path.exists(path):
                data = np.load(path, allow_pickle=True)
                self._W_s = data['W_sensory']  # (75, 300)
                self._P = data['P_motor']      # (300, 75)
                print(f"[Brain] 映射模式已启用: W_s={self._W_s.shape}, P={self._P.shape}")
        except Exception as e:
            print(f"[Brain] 映射矩阵未找到, 使用原始模式: {e}")
    
    def _load_qa_index(self):
        """加载QA检索索引: 对训练数据的问题编码，建立向量检索"""
        import os, json
        try:
            qa_path = os.path.join(os.path.dirname(__file__), '..', '..', 'data', 'qa_training.json')
            if not os.path.exists(qa_path):
                print(f"[Brain] QA训练数据未找到: {qa_path}")
                return
            with open(qa_path, 'r') as f:
                qa_data = json.load(f)
            from .semantic_encoder import get_encoder
            encoder = get_encoder()
            vecs, answers = [], []
            for item in qa_data:
                q = item.get('q') or item.get('question', '')
                a = item.get('a') or item.get('answer', '')
                if not q or not a:
                    continue
                v = encoder.encode(q)
                if v is not None and np.linalg.norm(v) > 0:
                    vecs.append(v[:self.dim])
                    answers.append(a)
            if vecs:
                self._qa_vecs = np.array(vecs, dtype=np.float32)
                self._qa_answers = answers
                print(f"[Brain] QA索引已加载: {len(answers)}条, vecs={self._qa_vecs.shape}")
        except Exception as e:
            print(f"[Brain] QA索引加载失败: {e}")
    
    def _qa_retrieve(self, input_vec, threshold=0.75):
        """用输入向量检索最匹配的QA对"""
        if self._qa_vecs is None or len(self._qa_answers) == 0:
            return None, 0.0
        v = input_vec[:self.dim]
        v_norm = np.linalg.norm(v)
        if v_norm < 1e-8:
            return None, 0.0
        # 批量cosine
        norms = np.linalg.norm(self._qa_vecs, axis=1, keepdims=True)
        norms = np.maximum(norms, 1e-8)
        sims = (self._qa_vecs @ v) / (norms.ravel() * v_norm)
        best_idx = int(np.argmax(sims))
        best_sim = float(sims[best_idx])
        if best_sim >= threshold:
            return self._qa_answers[best_idx], best_sim
        return None, best_sim
    
    def forward(self, input_vec) -> Dict[str, np.ndarray]:
        """
        前向传播 — 6区协作推理
        
        映射模式: 输入75维→W_s(75×300)→sensory→AreaNetwork全链路→motor(300)→P(300×75)→输出75维
        原始模式: 输入75维→AreaNetwork全链路→motor(300)
        """
        if isinstance(input_vec, str):
            return self.chat(input_vec)
        
        x = np.asarray(input_vec, dtype=np.float32).ravel()[:self.dim]
        if len(x) < self.dim:
            x = np.pad(x, (0, self.dim - len(x)))
        
        # AreaNetwork原始forward
        output, meta = self._net.forward(x)
        self._forward_count += 1
        
        # 从meta提取各区输出
        result = {}
        area_outputs = meta.get('area_outputs', {})
        for area_type, area_data in area_outputs.items():
            key = area_type.value if hasattr(area_type, 'value') else str(area_type)
            result[key] = area_data
        
        # 映射模式: 用W_s覆盖sensory, 用P投影motor
        if self._W_s is not None:
            # W_s直接映射: 输入→sensory(跳过随机sensory层)
            result['sensory'] = (x @ self._W_s).astype(np.float32)
        
        # 确保有motor输出
        if 'motor' not in result:
            result['motor'] = output
        result['_final'] = output
        
        return result
    
    def chat(self, text: str) -> Optional[Dict]:
        """
        文字输入 — 语义编码为75维向量后forward
        
        Args:
            text: 输入文字
            
        Returns:
            {'text': 回复文字, 'confidence': 置信度, 'areas': 区激活信息}
        """
        # 语义编码（vecs75查表，替代hash）
        from .semantic_encoder import get_encoder
        encoder = get_encoder()
        vec = encoder.encode(text)
        
        # QA检索: 编码向量匹配训练数据
        # cos>=0.75精确匹配, 0.55~0.75模糊匹配(标注低置信度)
        qa_answer, qa_sim = self._qa_retrieve(vec, threshold=0.55)
        if qa_answer:
            mode = 'qa_brain' if qa_sim >= 0.75 else 'qa_brain_fuzzy'
            return {
                'text': qa_answer,
                'confidence': float(qa_sim),
                'areas': ['qa_retrieve'],
                'decoded_words': [(qa_answer[:8], qa_sim)],
                'mode': mode,
            }
        
        result = self.forward(vec)
        
        # 从多个区域信号综合计算置信度
        confidences = []
        for key in ['sensory', 'association', 'prefrontal', 'motor']:
            arr = result.get(key)
            if isinstance(arr, np.ndarray) and arr.size > 0:
                confidences.append(float(np.max(np.abs(arr))))
        
        confidence = max(confidences) if confidences else 0.0
        
        # 语言解码层: motor 300维 → 自回归生成句子
        motor_vec = result.get('motor', result.get('_final', np.zeros(1)))
        if self._decoder is not None and motor_vec.size >= 300:
            generated = self._decoder.decode(motor_vec[:300], top_k=10)
            if generated and len(generated) > 1:
                return {
                    'text': generated,
                    'confidence': confidence,
                    'areas': list(result.keys()),
                    'decoded_words': [(generated[:8], confidence)],
                    'mode': 'language_decoder',
                }
        
        # 回退: 最近邻词解码(未训练时)
        if self._W_s is not None:
            sensory = result.get('sensory', np.zeros(1))
            decode_vec = sensory[:75] if sensory.size >= 75 else motor_vec[:min(75, motor_vec.size)]
        elif self._P is not None and motor_vec.size >= 300:
            decode_vec = (motor_vec[:300] @ self._P).astype(np.float32)
        else:
            decode_vec = motor_vec
        
        top_words = encoder.decode_nearest(decode_vec, top_k=5)
        
        if top_words and top_words[0][1] > 0.3:
            decoded = ' '.join(w for w, s in top_words[:3] if s > 0.3)
        else:
            decoded = f'置信度:{confidence:.3f}'
        
        return {
            'text': f'[Brain] {decoded}',
            'confidence': confidence,
            'areas': list(result.keys()),
            'decoded_words': top_words[:5],
        }
    
    def stats(self) -> Dict:
        """返回统计信息"""
        area_info = {}
        total_params = 0
        for area_type, area in self._net.areas.items():
            name = area_type.value
            col_count = len(area.columns)
            mc_count = sum(len(c.micro_columns) for c in area.columns)
            # 用微柱的total_params属性统计
            params = 0
            for col in area.columns:
                for mc in col.micro_columns:
                    if hasattr(mc, 'total_params'):
                        params += mc.total_params
            total_params += params
            area_info[name] = {
                'area_name': name,
                'columns': col_count,
                'micro_columns': mc_count,
                'params': params,
            }
        
        return {
            'areas': len(self._net.areas),
            'total_params': total_params,
            'area_details': area_info,
            'forward_count': self._forward_count,
        }
    
    def get_area_weights(self, area_name: str) -> np.ndarray:
        """导出指定区权重 — 深度收集所有numpy属性(含_synaptic内部)"""
        area_type = AreaType(area_name)
        area = self._net.areas[area_type]
        weights_list = []
        for col in area.columns:
            for mc in col.micro_columns:
                # 收集mc自身的所有numpy属性(含下划线)
                mc._weight_attrs = []
                for attr_name in dir(mc):
                    val = getattr(mc, attr_name, None)
                    if isinstance(val, np.ndarray) and val.ndim >= 1 and val.size < 1_000_000:
                        weights_list.append(val.ravel())
                        mc._weight_attrs.append(attr_name)
                # 深入_synaptic收集
                syn = getattr(mc, '_synaptic', None)
                if syn is not None:
                    syn._weight_attrs = []
                    for attr_name in dir(syn):
                        val = getattr(syn, attr_name, None)
                        if isinstance(val, np.ndarray) and val.ndim >= 1 and val.size < 1_000_000:
                            weights_list.append(val.ravel())
                            syn._weight_attrs.append(attr_name)
        if weights_list:
            return np.concatenate(weights_list)
        return np.array([], dtype=np.float32)
    
    def set_area_weights(self, area_name: str, weights: np.ndarray):
        """导入指定区权重 — 按记录的属性顺序回填(含_synaptic内部)"""
        area_type = AreaType(area_name)
        area = self._net.areas[area_type]
        offset = 0
        for col in area.columns:
            for mc in col.micro_columns:
                # 回填mc自身属性
                for attr_name in getattr(mc, '_weight_attrs', []):
                    val = getattr(mc, attr_name, None)
                    if isinstance(val, np.ndarray) and val.ndim >= 1:
                        n = val.size
                        if offset + n <= weights.size:
                            setattr(mc, attr_name, weights[offset:offset+n].reshape(val.shape))
                            offset += n
                # 回填_synaptic属性
                syn = getattr(mc, '_synaptic', None)
                if syn is not None:
                    for attr_name in getattr(syn, '_weight_attrs', []):
                        val = getattr(syn, attr_name, None)
                        if isinstance(val, np.ndarray) and val.ndim >= 1:
                            n = val.size
                            if offset + n <= weights.size:
                                setattr(syn, attr_name, weights[offset:offset+n].reshape(val.shape))
                                offset += n
    
    def train_decoder(self, qa_data: list, epochs: int = 10, lr: float = 0.01) -> Dict:
        """训练语言解码层 — 用QA数据teacher forcing训练W_out映射
        
        Args:
            qa_data: [{'q': 问题, 'a': 答案}, ...]
            epochs: 训练轮数
            lr: 学习率
        """
        if self._decoder is None:
            return {'error': '解码层未初始化'}
        from .semantic_encoder import get_encoder
        encoder = get_encoder()
        total_loss = 0.0
        n_trained = 0
        for epoch in range(epochs):
            epoch_loss = 0.0
            for item in qa_data:
                q = item.get('q') or item.get('question', '')
                a = item.get('a') or item.get('answer', '')
                if not q or not a:
                    continue
                # 问题编码 → forward获取motor输出
                vec = encoder.encode(q)
                if vec is None:
                    continue
                result = self.forward(vec)
                motor_vec = result.get('motor', np.zeros(self.config.get('motor_output_dim', 300)))
                # 维度适配: 截断或填充到decoder的motor_dim
                dec_motor_dim = self._decoder.motor_dim
                if motor_vec.size < dec_motor_dim:
                    motor_vec = np.pad(motor_vec, (0, dec_motor_dim - motor_vec.size))
                elif motor_vec.size > dec_motor_dim:
                    motor_vec = motor_vec[:dec_motor_dim]
                # 训练一步: motor向量 + 目标答案
                result = self._decoder.train_step(motor_vec, a, lr=lr)
                epoch_loss += result['loss'] if isinstance(result, dict) else result
                n_trained += 1
            total_loss = epoch_loss
        # 保存权重
        try:
            dec_path = os.path.join(os.path.dirname(__file__), '..', '..', 'models', 'decoder_weights.npz')
            self._decoder.save(dec_path)
        except Exception as e:
            print(f"[Brain] 解码层权重保存失败: {e}")
        return {'epochs': epochs, 'n_trained': n_trained, 'final_loss': total_loss}

    def train(self, texts: list, epochs: int = 1, lr: float = 0.01):
        """在线训练 — 赫布学习（训练时临时解冻）
        
        关键：必须forward向量而非字符串，否则走chat()分支绕过AreaNetwork
        """
        import time as _t
        from .semantic_encoder import get_encoder
        encoder = get_encoder()
        self._net.unfreeze()
        try:
            for epoch in range(epochs):
                for i, text in enumerate(texts):
                    t0 = _t.time()
                    vec = encoder.encode(text) if isinstance(text, str) else text
                    if vec is None:
                        continue
                    t1 = _t.time()
                    self.forward(vec)
                    t2 = _t.time()
                    self._net.learn(force=True)
                    t3 = _t.time()
                    print(f"  train[{i}] encode={t1-t0:.1f}s forward={t2-t1:.1f}s learn={t3-t2:.1f}s", flush=True)
        finally:
            self._net.freeze()
    
    def fedavg(self, area_name: str, incoming: np.ndarray, node_id: str = '') -> Dict:
        """联邦平均聚合"""
        local = self.get_area_weights(area_name)
        if local.size == 0:
            self.set_area_weights(area_name, incoming)
            return {'action': 'adopt', 'node': node_id}
        # 简单平均
        min_len = min(local.size, incoming.size)
        avg = (local[:min_len] + incoming[:min_len]) / 2
        # 补齐
        if local.size > min_len:
            avg = np.concatenate([avg, local[min_len:]])
        elif incoming.size > min_len:
            avg = np.concatenate([avg, incoming[min_len:]])
        self.set_area_weights(area_name, avg.astype(np.float32))
        return {'action': 'avg', 'local_size': local.size, 'incoming_size': incoming.size}
    
    @property
    def areas(self) -> Dict[str, object]:
        """兼容: 返回 {区名str: FunctionalArea}"""
        return {at.value: area for at, area in self._net.areas.items()}
    
    def get_stats(self) -> Dict:
        """兼容旧接口"""
        return self.stats()
    
    def save(self, path: str):
        """保存模型"""
        all_weights = {}
        for area_type, area in self._net.areas.items():
            name = area_type.value
            all_weights[name] = self.get_area_weights(name)
        np.savez(path, **all_weights)
    
    def load(self, path: str):
        """加载模型"""
        data = np.load(path, allow_pickle=True)
        for name in data.files:
            try:
                self.set_area_weights(name, data[name])
            except (ValueError, KeyError):
                pass