src/core/synaptic/l1_sparse.py

"""
L1 稀疏突触微柱
随机选择5-10%神经元对建立突触连接
支持赫布学习: ΔW = η·(pre·post - λ·W)

特点: 模拟生物大脑的节能模式
参数: ~1K/微柱
适用: 感觉/运动/丘脑区
"""
import numpy as np
from typing import Dict, Optional


class SparseSynapticMicroColumn:
    """L1稀疏突触微柱 - 神经元间有稀疏突触连接"""
    
    # 神经元类型比例（按功能柱类型）
    NEURON_RATIOS = {
        'sensory':  {'E': 0.75, 'I': 0.20, 'M': 0.05},
        'memory':   {'E': 0.80, 'I': 0.15, 'M': 0.05},
        'detector': {'E': 0.70, 'I': 0.20, 'M': 0.10},
        'integrator': {'E': 0.70, 'I': 0.15, 'M': 0.15},
        'selector': {'E': 0.70, 'I': 0.20, 'M': 0.10},
        'motor':    {'E': 0.80, 'I': 0.15, 'M': 0.05},
        'modulator': {'E': 0.45, 'I': 0.25, 'M': 0.30},
    }
    
    def __init__(self, column_type: str = 'sensory', num_neurons: int = 100,
                 sparsity: float = 0.08, learning_rate: float = 0.005,
                 decay_rate: float = 0.001,
                 receptive_field_size: int = None,
                 receptive_field_offset: int = 0):
        """
        Args:
            column_type: 功能类型
            num_neurons: 神经元数量
            sparsity: 突触连接稀疏度(5-10%)
            learning_rate: 赫布学习率
            decay_rate: 权重衰减率
            receptive_field_size: 感受野大小(None=全连接, int=只连接部分输入维度)
            receptive_field_offset: 感受野起始偏移(分块式: 每个mc负责不同输入区域)
        """
        self.column_type = column_type
        self.num_neurons = num_neurons
        self.sparsity = sparsity
        self.learning_rate = learning_rate
        self.decay_rate = decay_rate
        self.name = f"SparseSynaptic-{column_type}"
        
        # 神经元分组
        ratios = self.NEURON_RATIOS.get(column_type, self.NEURON_RATIOS['sensory'])
        self.n_e = int(num_neurons * ratios['E'])
        self.n_i = int(num_neurons * ratios['I'])
        self.n_m = num_neurons - self.n_e - self.n_i
        
        # 神经元状态
        self.membrane = np.zeros(num_neurons, dtype=np.float32)
        self.threshold = np.full(num_neurons, 1.0, dtype=np.float32)
        self.refractory = np.zeros(num_neurons, dtype=np.float32)
        
        # === 稀疏突触权重矩阵 ===
        # 4类突触连接: E→E, E→I, I→E, M→E
        self.W_ee = self._create_sparse(self.n_e, self.n_e)
        self.W_ei = self._create_sparse(self.n_e, self.n_i)
        self.W_ie = self._create_sparse(self.n_i, self.n_e)
        self.W_me = self._create_sparse(self.n_m, self.n_e)
        
        # 输入投影矩阵（外部输入→E神经元）— 支持感受野
        self.receptive_field_size = receptive_field_size
        self.receptive_field_offset = receptive_field_offset
        self.W_input = np.random.randn(self.n_e, num_neurons).astype(np.float32) * 0.1
        
        # 应用感受野: 随机选择rf_size个输入维度(列)
        # 注意: W_input shape=(n_e, num_neurons), 实际输入维度=input_dim(≤num_neurons)
        if receptive_field_size is not None and receptive_field_size < num_neurons:
            # 用offset做种子，确保可复现但每个mc不同
            rng = np.random.RandomState(receptive_field_offset + 42)
            # 只在真实输入维度范围内选择(避免选到padding区域)
            effective_input_dim = min(receptive_field_size * 7, num_neurons)  # 覆盖约7倍RF大小
            rf_indices = sorted(rng.choice(effective_input_dim, receptive_field_size, replace=False))
            # W_input shape = (n_e, num_neurons), 列=输入维度
            col_mask = np.zeros_like(self.W_input)
            col_mask[:, rf_indices] = 1
            self.W_input = self.W_input * col_mask
        
        # 学习统计
        self._forward_count = 0
        self._hebb_updates = 0
        
    def _create_sparse(self, n_pre: int, n_post: int) -> np.ndarray:
        """创建稀疏突触权重矩阵"""
        n_synapses = max(1, int(n_pre * n_post * self.sparsity))
        W = np.zeros((n_pre, n_post), dtype=np.float32)
        
        # 随机选择突触位置
        indices = np.random.choice(n_pre * n_post, n_synapses, replace=False)
        rows, cols = np.divmod(indices, n_post)
        # 初始权重: 兴奋性正, 抑制性负
        W[rows, cols] = np.random.randn(n_synapses).astype(np.float32) * 0.1
        
        # 记录哪些位置有突触（学习时只更新这些位置）
        mask = np.zeros_like(W, dtype=bool)
        mask[rows, cols] = True
        self._mask_ee = None  # 稍后设置
        W[~mask] = 0  # 确保非突触位置为零
        
        return W
    
    def _create_sparse_with_mask(self, n_pre: int, n_post: int) -> tuple:
        """创建稀疏矩阵+掩码"""
        n_synapses = max(1, int(n_pre * n_post * self.sparsity))
        W = np.zeros((n_pre, n_post), dtype=np.float32)
        
        indices = np.random.choice(n_pre * n_post, n_synapses, replace=False)
        rows, cols = np.divmod(indices, n_post)
        W[rows, cols] = np.random.randn(n_synapses).astype(np.float32) * 0.1
        
        mask = np.zeros((n_pre, n_post), dtype=bool)
        mask[rows, cols] = True
        
        return W, mask
    
    def forward(self, inputs: np.ndarray, learn: bool = False) -> np.ndarray:
        """前向传播
        
        Args:
            inputs: 外部输入信号
            learn: 是否执行赫布学习（默认False，由learn()显式调用）
        """
        frozen = getattr(self, '_frozen', False)
        if frozen:
            learn = False
        
        x = np.asarray(inputs, dtype=np.float32).flatten()
        if len(x) < self.num_neurons:
            x = np.pad(x, (0, self.num_neurons - len(x)))
        elif len(x) > self.num_neurons:
            x = x[:self.num_neurons]
        
        # Step 1: 外部输入 → E神经元
        e_input = self.W_input @ x
        
        # Step 2: 突触传播（2个时间步）
        # 时间步1: E→I, E→E
        i_input = self.W_ee.T[:self.n_i] @ e_input if self.n_i > 0 else np.zeros(0)
        e_lateral = self.W_ee @ e_input
        
        # E神经元激活 — 使用ReLU保留稀疏性和分化
        e_activation = np.maximum(0, e_input + e_lateral)
        
        # I神经元激活
        i_activation = np.maximum(0, i_input) if self.n_i > 0 else np.zeros(0)
        
        # 时间步2: I→E抑制, M→E调节
        e_inhibition = self.W_ie.T @ i_activation if self.n_i > 0 else np.zeros(self.n_e)
        e_modulation = self.W_me.T @ np.tanh(x[:self.n_m]) if self.n_m > 0 else np.zeros(self.n_e)
        
        # 最终E输出 — 单层激活（去掉第二层tanh压缩）
        e_output = np.maximum(0, e_activation - e_inhibition + e_modulation)
        
        # Winner-Take-All竞争: 只保留top-k激活，其余置零
        k = max(1, int(self.n_e * 0.3))  # 保留前30%
        if k < self.n_e:
            top_k_idx = np.argpartition(np.abs(e_output), -k)[-k:]
            mask = np.zeros(self.n_e, dtype=np.float32)
            mask[top_k_idx] = 1.0
            e_output = e_output * mask
        
        # 输出标准化: 归一化到单位球面，保留方向信息
        out_norm = np.linalg.norm(e_output)
        if out_norm > 1e-6:
            e_output = e_output / out_norm
        
        # Step 3: 赫布学习 — 由learn()显式调用，forward中不自动学习
        # (原设计: if learn: self._hebbian_update(...))
        # 改为: forward只做推理，learn()负责学习，避免forward隐式修改权重
        
        # 缓存激活用于外部训练
        self._last_pre = e_input.copy()
        self._last_post = e_output.copy()
        self._last_i_act = i_activation.copy() if self.n_i > 0 else np.zeros(0)
        
        # Step 4: 更新膜电位
        self.membrane[:self.n_e] = e_output
        if self.n_i > 0:
            self.membrane[self.n_e:self.n_e+self.n_i] = i_activation
        if self.n_m > 0:
            self.membrane[self.n_e+self.n_i:] = np.tanh(x[:self.n_m])
        
        # 不应期衰减
        self.refractory = np.maximum(0, self.refractory - 1)
        
        self._forward_count += 1
        return self.membrane.copy()
    
    def _hebbian_update(self, pre: np.ndarray, post: np.ndarray,
                        i_activation: np.ndarray):
        """改进赫布学习: Oja规则 + 去相关 + 行级L2归一化
        
        核心改进:
        1. Oja规则 ΔW = η·(pre·post - post²·W) 基础学习
        2. 反赫布去相关: 推开与最近pattern的相似度，实现分化
        3. 行级L2归一化防止权重膨胀
        """
        lr = self.learning_rate
        
        # E→E突触: Oja规则 (ΔW_ij = η·(x_i·y_j - y_j²·W_ij))
        n_pre = min(len(pre), self.W_ee.shape[0])
        n_post = min(len(post), self.W_ee.shape[1])
        
        pre_slice = pre[:n_pre]
        post_slice = post[:n_post]
        
        # Oja更新
        post_sq = post_slice ** 2
        oja_term = np.outer(np.ones(n_pre), post_sq) * self.W_ee[:n_pre, :n_post]
        delta = lr * (np.outer(pre_slice, post_slice) - oja_term)
        
        # 反赫布去相关: 如果输出与最近历史太相似，推开
        if not hasattr(self, '_output_history'):
            self._output_history = []
        self._output_history.append(post_slice.copy())
        if len(self._output_history) > 10:
            self._output_history = self._output_history[-10:]
        
        if len(self._output_history) >= 2:
            # 计算当前输出与历史均值的相似度
            mean_post = np.mean(self._output_history[:-1], axis=0)
            cos_sim = np.dot(post_slice, mean_post) / (np.linalg.norm(post_slice) * np.linalg.norm(mean_post) + 1e-8)
            # 相似度越高，去相关力越强
            decorr_strength = lr * 0.5 * max(0, cos_sim)
            if decorr_strength > 0:
                # 推开: 减弱与历史均值方向的连接
                decorr_delta = decorr_strength * np.outer(pre_slice, mean_post[:n_post])
                delta -= decorr_delta
        
        # 梯度裁剪
        max_delta = 0.1 * np.abs(self.W_ee[:n_pre, :n_post]) + 1e-6
        delta = np.clip(delta, -max_delta, max_delta)
        
        self.W_ee[:n_pre, :n_post] += delta
        
        # 行级L2归一化
        row_norms = np.linalg.norm(self.W_ee, axis=1, keepdims=True)
        row_norms = np.maximum(row_norms, 1e-8)
        self.W_ee = self.W_ee / row_norms * np.clip(row_norms, 0, 1.5)
        
        # E→I突触: 标准赫布+衰减
        if self.n_i > 0 and self.W_ei.size > 0:
            self.W_ei += lr * (np.outer(pre[:self.W_ei.shape[0]], i_activation[:self.W_ei.shape[1]]) - self.decay_rate * self.W_ei)
            self.W_ei = np.clip(self.W_ei, -1.5, 1.5)
        
        # I→E突触: 标准赫布+衰减
        if self.n_i > 0 and self.W_ie.size > 0:
            self.W_ie += lr * (np.outer(i_activation[:self.W_ie.shape[0]], post[:self.W_ie.shape[1]]) - self.decay_rate * self.W_ie)
            self.W_ie = np.clip(self.W_ie, -1.5, 1.5)
        
        self._hebb_updates += 1
    
    def get_synapse_count(self) -> int:
        """获取突触总数"""
        count = 0
        for W in [self.W_ee, self.W_ei, self.W_ie, self.W_me]:
            count += np.count_nonzero(W)
        return count
    
    def get_param_count(self) -> int:
        """获取可学习参数总数"""
        params = 0
        for W in [self.W_ee, self.W_ei, self.W_ie, self.W_me, self.W_input]:
            params += W.size
        params += self.threshold.size
        return params
    
    def get_config(self) -> Dict:
        return {
            'type': self.name,
            'tier': 'L1',
            'column_type': self.column_type,
            'num_neurons': self.num_neurons,
            'neurons': {'E': self.n_e, 'I': self.n_i, 'M': self.n_m},
            'sparsity': self.sparsity,
            'synapse_count': self.get_synapse_count(),
            'param_count': self.get_param_count(),
            'learning_rate': self.learning_rate,
            'forward_count': self._forward_count,
            'hebb_updates': self._hebb_updates,
        }
    
    def reset(self):
        """重置状态（保留突触权重）"""
        self.membrane = np.zeros(self.num_neurons, dtype=np.float32)
        self.refractory = np.zeros(self.num_neurons, dtype=np.float32)
    
    def learn(self):
        """执行赫布学习（供v3微柱调用）
        
        改进：使用真实缓存的输入/输出信号，而非近似值
        """
        # 优先使用真实缓存的输入输出信号
        pre = getattr(self, '_last_pre', None)
        post = getattr(self, '_last_post', None)
        
        if pre is not None and post is not None and len(pre) > 0 and len(post) > 0:
            # 使用真实信号
            e_input = pre[:self.n_e] if len(pre) >= self.n_e else pre
            e_output = post[:self.n_e] if len(post) >= self.n_e else post
        else:
            # 回退：使用当前膜电位
            if self._forward_count > 0:
                e_output = self.membrane[:self.n_e]
                e_input = self.W_input @ np.ones(self.num_neurons, dtype=np.float32) * 0.5
            else:
                return
        
        i_activation = self.membrane[self.n_e:self.n_e+self.n_i] if self.n_i > 0 else np.zeros(0)
        self._hebbian_update(e_input, e_output, i_activation)
    
    @property
    def total_params(self) -> int:
        """可学习参数数量（突触权重）"""
        return sum(w.size for w in [
            self.W_ee, self.W_ei, self.W_ie, self.W_me, self.W_input
        ])