src/core/micro_columns/association_cache.py

"""
AssociationCache: 联合区缓存模块
功能：临时存储 + 信息路由 + 多源融合
"""
import numpy as np
from typing import Dict, List, Tuple, Optional


class AssociationCache:
    """
    联合区缓存 - 信息传输枢纽
    
    功能：
    1. 临时存储多区输入（感觉区、记忆区、丘脑反馈）
    2. 信息路由：根据目标区域选择输出
    3. 多源融合：加权合并多路输入
    """
    
    def __init__(
        self,
        capacity: int = 10,      # 缓存容量（多少个时间步）
        decay_rate: float = 0.9, # 自然衰减率
        fusion_weights: Dict[str, float] = None  # 多源融合权重
    ):
        self.capacity = capacity
        self.decay_rate = decay_rate
        
        # 默认融合权重
        self.fusion_weights = fusion_weights or {
            'sensory': 0.4,    # 感觉输入权重最高（实时性）
            'memory': 0.3,     # 记忆输入
            'thalamus': 0.2,   # 丘脑反馈
            'prefrontal': 0.1  # 前额区残差
        }
        
        # 缓存存储：key=区域名, value=(时间步, 数据)
        self._cache: Dict[str, Tuple[int, np.ndarray]] = {}
        self._cache_order: List[str] = []  # 记录写入顺序（LRU）
        self._time_step = 0
        
        # 路由表：key=目标区域, value=源区域列表
        self._routing_table: Dict[str, List[str]] = {
            'prefrontal': ['sensory', 'memory', 'thalamus'],
            'motor': ['prefrontal'],
            'thalamus': ['prefrontal', 'motor'],
            'memory': ['sensory', 'prefrontal']
        }
        
        # 统计
        self._stats = {
            'cache_hits': 0,
            'cache_misses': 0,
            'routing_count': 0,
            'fusion_count': 0
        }
    
    def store(self, source: str, data: np.ndarray) -> None:
        """
        存储数据到缓存
        
        Args:
            source: 来源区域 ('sensory', 'memory', 'thalamus', 'prefrontal')
            data: 要存储的数据
        """
        self._time_step += 1
        
        # LRU淘汰：如果超过容量，删除最旧的
        if len(self._cache) >= self.capacity and source not in self._cache:
            oldest = self._cache_order[0] if self._cache_order else None
            if oldest:
                del self._cache[oldest]
                self._cache_order.remove(oldest)
        
        # 存储新数据
        self._cache[source] = (self._time_step, data.copy())
        
        # 更新LRU顺序
        if source in self._cache_order:
            self._cache_order.remove(source)
        self._cache_order.append(source)
    
    def retrieve(self, source: str) -> Optional[np.ndarray]:
        """
        从缓存检索数据
        
        Args:
            source: 来源区域
            
        Returns:
            缓存的数据，如果不存在返回None
        """
        if source in self._cache:
            self._stats['cache_hits'] += 1
            return self._cache[source][1].copy()
        else:
            self._stats['cache_misses'] += 1
            return None
    
    def route(self, target: str) -> Optional[np.ndarray]:
        """
        路由：根据目标区域从缓存中选取数据
        
        Args:
            target: 目标区域
            
        Returns:
            路由的数据（如果有多源则融合）
        """
        self._stats['routing_count'] += 1
        
        # 查找目标区域的源区域
        source_regions = self._routing_table.get(target, [])
        
        # 收集可用的源数据
        available_sources = []
        for src in source_regions:
            data = self.retrieve(src)
            if data is not None:
                available_sources.append((src, data))
        
        if not available_sources:
            return None
        
        # 如果只有一个源，直接返回
        if len(available_sources) == 1:
            return available_sources[0][1]
        
        # 多源融合
        return self._fuse_sources(available_sources)
    
    def _fuse_sources(self, sources: List[Tuple[str, np.ndarray]]) -> np.ndarray:
        """
        多源融合
        
        Args:
            sources: [(源区域, 数据), ...]
            
        Returns:
            融合后的数据
        """
        self._stats['fusion_count'] += 1
        
        fused = np.zeros_like(sources[0][1])
        total_weight = 0.0
        
        for src, data in sources:
            weight = self.fusion_weights.get(src, 0.25)
            fused += weight * data
            total_weight += weight
        
        # 归一化
        if total_weight > 0:
            fused = fused / total_weight
        
        return fused
    
    def decay_all(self) -> None:
        """对所有缓存应用衰减"""
        for key in list(self._cache.keys()):
            time_step, data = self._cache[key]
            age = self._time_step - time_step
            decay = self.decay_rate ** age
            self._cache[key] = (time_step, data * decay)
    
    def clear(self) -> None:
        """清空缓存"""
        self._cache.clear()
        self._cache_order.clear()
    
    def get_stats(self) -> Dict:
        """获取统计信息"""
        total = self._stats['cache_hits'] + self._stats['cache_misses']
        hit_rate = self._stats['cache_hits'] / total if total > 0 else 0.0
        
        return {
            'cache_size': len(self._cache),
            'cache_capacity': self.capacity,
            'cache_hits': self._stats['cache_hits'],
            'cache_misses': self._stats['cache_misses'],
            'hit_rate': hit_rate,
            'routing_count': self._stats['routing_count'],
            'fusion_count': self._stats['fusion_count']
        }


class MultiRegionRouter:
    """
    多区域路由器 - 控制信息流向
    
    功能：
    1. 决定信息从哪个区域流向哪个区域
    2. 条件路由：根据当前状态选择不同路径
    3. 广播：一对多信息分发
    """
    
    def __init__(self):
        # 路由策略：key=源区域, value={目标区域: 条件函数}
        self._routes: Dict[str, Dict[str, callable]] = {}
        
        # 当前激活路径
        self._active_routes: List[Tuple[str, str]] = []
        
        # 路由历史
        self._route_history: List[Dict] = []
    
    def add_route(
        self,
        source: str,
        target: str,
        condition: callable = None
    ) -> None:
        """
        添加路由规则
        
        Args:
            source: 源区域
            target: 目标区域
            condition: 条件函数（可选），接收当前状态返回bool
        """
        if source not in self._routes:
            self._routes[source] = {}
        self._routes[source][target] = condition
    
    def route(
        self,
        source: str,
        data: np.ndarray,
        state: Dict = None
    ) -> List[Tuple[str, np.ndarray]]:
        """
        执行路由
        
        Args:
            source: 源区域
            data: 要路由的数据
            state: 当前状态（用于条件判断）
            
        Returns:
            [(目标区域, 数据), ...]
        """
        if source not in self._routes:
            return [(source, data)]  # 直通
        
        state = state or {}
        results = []
        
        for target, condition in self._routes[source].items():
            # 检查条件
            if condition is None or condition(state):
                results.append((target, data.copy()))
                self._active_routes.append((source, target))
        
        # 记录路由历史
        self._route_history.append({
            'source': source,
            'targets': [t for t, _ in results],
            'state_snapshot': state.copy() if state else {}
        })
        
        # 保持历史在100条以内
        if len(self._route_history) > 100:
            self._route_history = self._route_history[-100:]
        
        return results if results else [(source, data)]
    
    def broadcast(
        self,
        data: np.ndarray,
        targets: List[str]
    ) -> List[Tuple[str, np.ndarray]]:
        """
        广播到多个目标
        
        Args:
            data: 要广播的数据
            targets: 目标区域列表
            
        Returns:
            [(目标区域, 数据), ...]
        """
        return [(t, data.copy()) for t in targets]
    
    def get_active_routes(self) -> List[Tuple[str, str]]:
        """获取当前激活的路由"""
        return self._active_routes.copy()
    
    def clear_routes(self) -> None:
        """清空激活路由"""
        self._active_routes.clear()