"""
L2 分层突触微柱（推荐档位）
E层 → I层 → E层 → M层 按功能分层连接
支持赫布学习: ΔW = η·(pre·post - λ·W)

特点: 模拟大脑皮层结构，可解释性强
参数: ~11K/微柱
适用: 记忆/联合区
"""
import numpy as np
from typing import Dict, Optional


class LayeredSynapticMicroColumn:
    """L2分层突触微柱 - E→I→E→M四层结构"""
    
    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 = 'memory', num_neurons: int = 100,
                 learning_rate: float = 0.01, decay_rate: float = 0.001):
        self.column_type = column_type
        self.num_neurons = num_neurons
        self.learning_rate = learning_rate
        self.decay_rate = decay_rate
        self.name = f"LayeredSynaptic-{column_type}"
        
        ratios = self.NEURON_RATIOS.get(column_type, self.NEURON_RATIOS['memory'])
        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)
        
        # === 分层突触权重 ===
        # Layer 1: 输入 → E1层（初级兴奋）
        self.W_input_e1 = np.random.randn(self.n_e, num_neurons).astype(np.float32) * 0.1
        
        # Layer 2: E1 → I层（横向抑制）
        self.W_e1_i = np.random.randn(self.n_i, self.n_e).astype(np.float32) * 0.1
        
        # Layer 3: I → E2层（抑制调节后的兴奋）
        self.W_i_e2 = np.random.randn(self.n_e, self.n_i).astype(np.float32) * 0.1
        # E1残差连接到E2
        self.W_e1_e2 = np.random.randn(self.n_e, self.n_e).astype(np.float32) * 0.05
        
        # Layer 4: E2 → M层（输出调制）
        self.W_e2_m = np.random.randn(self.n_m, self.n_e).astype(np.float32) * 0.1
        
        # M → E1 反馈连接（循环调节）
        self.W_m_e1 = np.random.randn(self.n_e, self.n_m).astype(np.float32) * 0.05
        
        # 学习统计
        self._forward_count = 0
        self._hebb_updates = 0
        
    def forward(self, inputs: np.ndarray, learn: bool = None) -> np.ndarray:
        """前向传播：E1→I→E2→M 四层处理
        
        Args:
            learn: 是否赫布学习。None=自动(非冻结时学习), False=不学习, True=强制学习
        """
        frozen = getattr(self, '_frozen', False)
        
        # learn参数自动推断: 冻结时不学习
        if learn is None:
            learn = not frozen
        
        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]
        
        # M层反馈（上一步的M状态影响当前E1）
        # 冻结模式: 用零反馈避免依赖历史状态
        if frozen or self._forward_count == 0:
            m_feedback = 0
        else:
            m_feedback = self.W_m_e1 @ self.membrane[self.n_e+self.n_i:] if self.n_m > 0 else 0
        
        # Layer 1: 输入 → E1
        e1_raw = self.W_input_e1 @ x + m_feedback
        e1 = np.tanh(e1_raw)
        
        # Layer 2: E1 → I（横向抑制）
        i_raw = self.W_e1_i @ e1
        i_activation = np.tanh(i_raw)
        
        # Layer 3: I → E2 + E1残差
        e2_raw = self.W_i_e2 @ i_activation + self.W_e1_e2 @ e1
        e2 = np.tanh(e2_raw)
        
        # Layer 4: E2 → M（输出调制）
        if self.n_m > 0:
            m_raw = self.W_e2_m @ e2
            m_activation = np.tanh(m_raw)
        else:
            m_activation = np.zeros(0)
        
        # 赫布学习
        if learn:
            self._hebbian_update(x, e1, i_activation, e2, m_activation)
        
        # 更新膜电位
        if not frozen:
            self.membrane[:self.n_e] = e2
            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:] = m_activation
            self._forward_count += 1
        
        return self.membrane.copy()
    
    def _hebbian_update(self, x, e1, i_act, e2, m_act):
        """赫布学习"""
        lr = self.learning_rate
        decay = self.decay_rate
        
        # E1→I: 输入-抑制关联
        if self.n_i > 0:
            self.W_e1_i += lr * (np.outer(i_act, e1) - decay * self.W_e1_i)
            self.W_e1_i = np.clip(self.W_e1_i, -2.0, 2.0)
        
        # I→E2: 抑制-兴奋关联
        if self.n_i > 0:
            self.W_i_e2 += lr * (np.outer(e2, i_act) - decay * self.W_i_e2)
            self.W_i_e2 = np.clip(self.W_i_e2, -2.0, 2.0)
        
        # E1→E2: 兴奋残差关联
        self.W_e1_e2 += lr * (np.outer(e2, e1) - decay * self.W_e1_e2)
        self.W_e1_e2 = np.clip(self.W_e1_e2, -2.0, 2.0)
        
        # E2→M: 输出-调制关联
        if self.n_m > 0 and len(m_act) > 0:
            self.W_e2_m += lr * (np.outer(m_act, e2) - decay * self.W_e2_m)
            self.W_e2_m = np.clip(self.W_e2_m, -2.0, 2.0)
        
        self._hebb_updates += 1
    
    def get_param_count(self) -> int:
        total = 0
        for W in [self.W_input_e1, self.W_e1_i, self.W_i_e2, 
                  self.W_e1_e2, self.W_e2_m, self.W_m_e1]:
            total += W.size
        total += self.threshold.size
        return total
    
    def get_config(self) -> Dict:
        return {
            'type': self.name,
            'tier': 'L2',
            'column_type': self.column_type,
            'num_neurons': self.num_neurons,
            'neurons': {'E': self.n_e, 'I': self.n_i, 'M': self.n_m},
            '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)
    
    def learn(self):
        """执行赫布学习（供v3微柱调用）"""
        if self._forward_count > 0:
            # 使用当前膜电位作为激活
            e2 = self.membrane[:self.n_e]
            e1 = self.membrane[:self.n_e] * 0.5
            i_activation = self.membrane[self.n_e:self.n_e+self.n_i] if self.n_i > 0 else np.zeros(0)
            m_activation = self.membrane[self.n_e+self.n_i:] if self.n_m > 0 else np.zeros(0)
            x = np.ones(self.num_neurons, dtype=np.float32) * 0.5
            self._hebbian_update(x, e1, i_activation, e2, m_activation)
    
    @property
    def total_params(self) -> int:
        """可学习参数数量（突触权重）"""
        return sum(w.size for w in [
            self.W_input_e1, self.W_e1_i, self.W_i_e2,
            self.W_e1_e2, self.W_e2_m, self.W_m_e1
        ])