"""
赫布学习协调器 - 跨微柱、跨功能柱的协同学习
实现"一起激活的神经元连接在一起"的赫布法则
"""
import numpy as np
from typing import Dict, List, Optional, Tuple
from enum import Enum


class HebbianMode(Enum):
    """赫布学习模式"""
    STANDARD = 'standard'           # 标准: ΔW = η * pre * post
    OJA = 'oja'                     # Oja法则: 加入权重复归一化
    BCM = 'bcm'                    # BCM理论: 稳定可塑性
    HOMEOSTATIC = 'homeostatic'    # 稳态: 抑制过活跃神经元


class HebbianLearningCoordinator:
    """
    赫布学习协调器 - 管理整个网络的协同学习
    
    功能:
    - 微柱内赫布学习（同一功能柱内微柱之间）
    - 微柱间赫布学习（不同功能柱之间）
    - 功能区之间赫布学习（跨区域联想）
    - 可塑性调节（根据任务动态调整学习率）
    """
    
    def __init__(
        self,
        learning_rate: float = 0.01,
        mode: HebbianMode = HebbianMode.OJA,
        decay_rate: float = 0.001,
        min_learning_rate: float = 0.001,
        max_learning_rate: float = 0.1
    ):
        self.base_learning_rate = learning_rate
        self.current_learning_rate = learning_rate
        self.mode = mode
        self.decay_rate = decay_rate
        self.min_lr = min_learning_rate
        self.max_lr = max_learning_rate
        
        # 微柱间连接权重（跨微柱学习）
        self._inter_column_weights: Dict[Tuple[str, str], np.ndarray] = {}
        
        # 激活历史（用于赫布计算）
        self._activation_history: Dict[str, List[np.ndarray]] = {}
        self._max_history = 20
        
        # 神经元活跃度统计（稳态调节）
        self._activity_stats: Dict[str, float] = {}  # 平均活跃度
        
        # 学习统计
        self._total_updates = 0
        self._coactivity_events = 0
        
    # ============ 核心赫布学习公式 ============
    
    def compute_hebbian_update(
        self,
        pre_activity: np.ndarray,
        post_activity: np.ndarray,
        current_weight: Optional[np.ndarray] = None,
        mode: Optional[HebbianMode] = None
    ) -> np.ndarray:
        """
        计算赫布学习权重更新
        
        ΔW = η * pre * post - λ * W  (标准)
        ΔW = η * pre * post - η * β * W * post²  (Oja)
        
        Args:
            pre_activity: 前突触激活
            post_activity: 后突触激活
            current_weight: 当前权重矩阵（如有）
            mode: 学习模式
            
        Returns:
            weight_update: 权重更新量
        """
        mode = mode or self.mode
        lr = self.current_learning_rate
        
        # 外积得到更新矩阵
        if pre_activity.ndim == 1:
            pre = pre_activity.reshape(-1, 1)
            post = post_activity.reshape(1, -1)
        else:
            pre = pre_activity
            post = post_activity
            
        hebbian_update = lr * np.dot(pre, post)
        
        if mode == HebbianMode.STANDARD and current_weight is not None:
            # 标准: 加上衰减项
            update = hebbian_update - lr * self.decay_rate * current_weight
            
        elif mode == HebbianMode.OJA and current_weight is not None:
            # Oja法则: 抑制性项与post²成正比
            post_squared = np.sum(post_activity ** 2)
            inhibition = lr * self.decay_rate * post_squared
            update = hebbian_update - inhibition * current_weight
            
        elif mode == HebbianMode.BCM:
            # BCM: 使用阈值φ(post)
            threshold = np.mean(post_activity)
            phi_post = post_activity * (post_activity - threshold)
            phi_post = phi_post.reshape(-1, 1)
            update = lr * (pre_activity.reshape(-1, 1) @ phi_post.T)
            # 衰减
            if current_weight is not None:
                update -= lr * self.decay_rate * current_weight
                
        elif mode == HebbianMode.HOMEOSTATIC:
            # 稳态: 抑制过活跃的
            activity_ratio = np.mean(post_activity) / (self._activity_stats.get('target', 0.5) + 1e-8)
            if activity_ratio > 1.0:
                # 过活跃，降低学习率
                self.current_learning_rate = max(self.min_lr, lr * 0.5)
            else:
                self.current_learning_rate = min(self.max_lr, lr * 1.2)
            update = hebbian_update
            if current_weight is not None:
                update -= lr * self.decay_rate * current_weight
        else:
            update = hebbian_update
            
        return update
    
    # ============ 微柱间学习 ============
    
    def register_micro_column(self, mc_name: str):
        """注册微柱以跟踪其激活"""
        if mc_name not in self._activation_history:
            self._activation_history[mc_name] = []
            self._activity_stats[mc_name] = 0.0
    
    def record_activation(self, mc_name: str, activation: np.ndarray):
        """记录微柱激活历史"""
        if mc_name not in self._activation_history:
            self.register_micro_column(mc_name)
            
        act = np.asarray(activation, dtype=np.float32).flatten()
        self._activation_history[mc_name].append(act.copy())
        
        # 保持历史长度
        if len(self._activation_history[mc_name]) > self._max_history:
            self._activation_history[mc_name].pop(0)
            
        # 更新活跃度统计
        self._activity_stats[mc_name] = 0.9 * self._activity_stats.get(mc_name, 0) + 0.1 * np.mean(np.abs(act))
    
    def learn_inter_column(
        self,
        source_mc: str,
        target_mc: str,
        source_activation: np.ndarray,
        target_activation: np.ndarray
    ) -> np.ndarray:
        """
        微柱间赫布学习
        
        Args:
            source_mc: 源微柱名
            target_mc: 目标微柱名
            source_activation: 源激活
            target_activation: 目标激活
            
        Returns:
            weight_update: 权重更新
        """
        key = (source_mc, target_mc)
        
        # 获取或初始化权重
        if key not in self._inter_column_weights:
            dim = min(len(source_activation), len(target_activation))
            self._inter_column_weights[key] = np.random.randn(dim, dim).astype(np.float32) * 0.01
            
        current_weights = self._inter_column_weights[key]
        
        # 调整维度
        src = source_activation.flatten()[:current_weights.shape[0]]
        tgt = target_activation.flatten()[:current_weights.shape[1]]
        
        # 计算赫布更新
        update = self.compute_hebbian_update(src, tgt, current_weights)
        
        # 应用更新
        self._inter_column_weights[key] = np.clip(
            current_weights + update,
            -2.0, 2.0
        )
        
        self._total_updates += 1
        return update
    
    def learn_coactivity(
        self,
        mc_pairs: List[Tuple[str, str]],
        activations: Dict[str, np.ndarray]
    ) -> int:
        """
        多微柱共活跃学习
        
        Args:
            mc_pairs: 微柱对列表 [(mc1, mc2), ...]
            activations: 微柱激活字典 {mc_name: activation}
            
        Returns:
            updated_pairs: 更新的对数
        """
        updated = 0
        for src, tgt in mc_pairs:
            if src in activations and tgt in activations:
                self.learn_inter_column(
                    src, tgt,
                    activations[src],
                    activations[tgt]
                )
                self._coactivity_events += 1
                updated += 1
                
        return updated
    
    # ============ 区域间学习 ============
    
    def learn_area_association(
        self,
        area_a_name: str,
        area_b_name: str,
        area_a_output: np.ndarray,
        area_b_output: np.ndarray
    ) -> np.ndarray:
        """
        功能区之间的联想学习
        
        两个区域同时活跃时，增强它们之间的联系
        """
        key = (f"area_{area_a_name}", f"area_{area_b_name}")
        
        if key not in self._inter_column_weights:
            dim_a = len(area_a_output)
            dim_b = len(area_b_output)
            self._inter_column_weights[key] = np.random.randn(dim_a, dim_b).astype(np.float32) * 0.01
            
        weights = self._inter_column_weights[key]
        
        # 确保维度匹配
        min_a = min(dim_a, weights.shape[0])
        min_b = min(dim_b, weights.shape[1])
        
        src = area_a_output.flatten()[:min_a]
        tgt = area_b_output.flatten()[:min_b]
        
        # 赫布更新
        update = self.compute_hebbian_update(src, tgt, weights[:min_a, :min_b])
        
        # 应用
        self._inter_column_weights[key][:min_a, :min_b] = np.clip(
            weights[:min_a, :min_b] + update,
            -2.0, 2.0
        )
        
        return update
    
    # ============ 可塑性调节 ============
    
    def adjust_learning_rate(self, task_difficulty: float, feedback_quality: float):
        """
        根据任务难度和反馈质量动态调整学习率
        
        Args:
            task_difficulty: 0-1, 任务难度
            feedback_quality: 0-1, 反馈质量（正确=1，错误=0）
        """
        # 难任务需要更大学习率
        difficulty_factor = 1.0 + task_difficulty
        
        # 高质量反馈可以大学习率
        feedback_factor = 0.5 + 0.5 * feedback_quality
        
        # 计算新学习率
        new_lr = self.base_learning_rate * difficulty_factor * feedback_factor
        self.current_learning_rate = np.clip(new_lr, self.min_lr, self.max_lr)
    
    def plasticity_regulation(
        self,
        recent_accuracy: float,
        stability_threshold: float = 0.9
    ) -> str:
        """
        可塑性调节 - 根据准确性动态调整
        
        Args:
            recent_accuracy: 近期准确率
            stability_threshold: 稳定阈值
            
        Returns:
            regulation_type: 'increase', 'decrease', 'maintain'
        """
        if recent_accuracy < stability_threshold - 0.1:
            # 准确性下降，增加可塑性
            self.current_learning_rate = min(
                self.max_lr,
                self.current_learning_rate * 1.5
            )
            return 'increase'
        elif recent_accuracy > stability_threshold:
            # 准确性很高，减少可塑性（巩固学习）
            self.current_learning_rate = max(
                self.min_lr,
                self.current_learning_rate * 0.8
            )
            return 'decrease'
        else:
            return 'maintain'
    
    # ============ 权重传播 ============
    
    def propagate_activity(
        self,
        source_mc: str,
        target_mc: str,
        activation: np.ndarray
    ) -> np.ndarray:
        """
        通过学习到的权重传播激活
        
        Args:
            source_mc: 源微柱
            target_mc: 目标微柱
            activation: 源激活
            
        Returns:
            propagated: 传播后的激活
        """
        key = (source_mc, target_mc)
        if key not in self._inter_column_weights:
            return activation
            
        weights = self._inter_column_weights[key]
        
        # 维度适配
        src = activation.flatten()[:weights.shape[0]]
        
        # 矩阵乘法
        propagated = np.dot(weights, src)
        
        return propagated
    
    # ============ 统计与查询 ============
    
    def get_inter_column_weights(self, source_mc: str, target_mc: str) -> Optional[np.ndarray]:
        """获取微柱间权重"""
        key = (source_mc, target_mc)
        return self._inter_column_weights.get(key)
    
    def get_learning_stats(self) -> Dict:
        """获取学习统计"""
        return {
            'total_updates': self._total_updates,
            'coactivity_events': self._coactivity_events,
            'current_learning_rate': self.current_learning_rate,
            'learning_mode': self.mode.value,
            'registered_micro_columns': len(self._activation_history),
            'inter_column_connections': len(self._inter_column_weights),
            'activity_stats': {k: float(v) for k, v in self._activity_stats.items()}
        }
    
    def reset_learning(self):
        """重置学习状态（保留权重）"""
        self._activation_history = {k: [] for k in self._activation_history}
        self._activity_stats = {k: 0.0 for k in self._activity_stats}
    
    def clear_weights(self):
        """清除所有学习到的权重"""
        self._inter_column_weights = {}
        self._total_updates = 0
        self._coactivity_events = 0