"""
ManagerMicroColumn: 管理微柱 - 调度+监控+拆分

核心职责:
1. 注册: 记录下级微柱的能力标签
2. 调度: 输入来了，路由到匹配的微柱
3. 监控: 跟踪各微柱的负载/记忆量
4. 拆分: 负载达上限时，分裂成两个更专精的微柱(有重合)

设计理念:
- 不是扩容，是分裂。像细胞分裂，越分越专精
- 初始少量微柱处理粗粒度任务，随学习积累自动拆分细化
- 拆分后有重合，保证过渡期不丢能力
"""

import numpy as np
from typing import Dict, List, Tuple, Optional


class ManagedUnit:
    """被管理的下级单元信息"""
    def __init__(self, unit_id: str, tags: List[str], 
                 capacity: int = 100):
        self.unit_id = unit_id
        self.tags = tags              # 能力标签，如 ['text','semantic']
        self.capacity = capacity      # 容量上限
        self.load = 0                 # 当前负载(记忆条数等)
        self.activation_count = 0     # 被激活次数
    
    @property
    def utilization(self) -> float:
        """利用率 0~1"""
        return self.load / self.capacity if self.capacity > 0 else 0.0


class ManagerMicroColumn:
    """
    管理微柱 - 多维度调度+自动拆分
    
    核心设计:
    - 每个管理微柱负责一个管理维度(如:频率/时序/语义)
    - 管理维度从学习中产生，初始1个维度，随数据积累分裂
    - 维度标签: 如 'frequency', 'temporal', 'semantic'
    - 业务微柱注册时带上该维度的标签值
    
    自身分裂:
    - 当管理微柱发现同一维度下路由冲突太多(多个不相关单元总被一起激活)
    - 说明这个维度粒度太粗，需要自身分裂成2个更细的维度
    - 例: 'feature'维度 → 分裂成 'edge_feature' + 'texture_feature'
    """
    
    SPLIT_THRESHOLD = 0.85
    # 维度分裂阈值: 路由冲突率超过此值，管理微柱自身分裂
    DIMENSION_SPLIT_THRESHOLD = 0.6
    
    def __init__(self, num_neurons: int = 64, 
                 dimension: str = 'default',
                 split_threshold: float = 0.85):
        self.num_neurons = num_neurons
        self.name = "Manager"
        self.function = f"调度+监控+拆分(维度:{dimension})"
        
        # 管理维度标签
        self.dimension = dimension
        
        self.split_threshold = split_threshold
        
        # 被管理的下级单元注册表
        self._registry: Dict[str, ManagedUnit] = {}
        
        # 路由冲突计数(同一次调度激活太多不相关单元)
        self._route_conflicts = 0
        self._total_dispatches = 0
        
        # 信号分类权重 (输入→标签空间)
        self._tag_dim = 32
        self._W_classify = np.random.randn(
            self._tag_dim, num_neurons
        ).astype(np.float32) * 0.1
        
        # 标签到单元的映射权重
        self._W_route = None
        
        # 拆分历史记录
        self._split_history: List[Dict] = []
        # 自身维度分裂历史
        self._dimension_split_history: List[Dict] = []
        
        self._dispatch_count = 0
    
    def forward(self, x: np.ndarray) -> np.ndarray:
        """前向传播 - 管理微柱直接透传输入（不做处理）"""
        # 管理微柱不直接处理信号，而是调度其他微柱
        # 此处透传，保持信号流
        return x
    
    def register(self, unit_id: str, tags: List[str], 
                 capacity: int = 100, load: int = 0):
        """注册下级单元
        
        Args:
            unit_id: 单元唯一ID
            tags: 能力标签列表，如 ['text','semantic']
            capacity: 容量上限
            load: 当前负载
        """
        self._registry[unit_id] = ManagedUnit(
            unit_id, tags, capacity
        )
        self._registry[unit_id].load = load
        self._rebuild_route_matrix()
    
    def unregister(self, unit_id: str):
        """注销下级单元(拆分后替换旧单元)"""
        if unit_id in self._registry:
            del self._registry[unit_id]
            self._rebuild_route_matrix()
    
    def update_load(self, unit_id: str, load: int):
        """更新下级单元负载"""
        if unit_id in self._registry:
            self._registry[unit_id].load = load
    
    def dispatch(self, input_signal: np.ndarray, 
                 input_tags: List[str] = None
                ) -> List[str]:
        """调度: 决定激活哪些下级单元
        
        Args:
            input_signal: 输入信号向量
            input_tags: 显式标签(如知道是文字输入)，None则自动推断
            
        Returns:
            激活的单元ID列表
        """
        if not self._registry:
            return []
        
        # 标签获取: 显式指定 或 自动推断
        if input_tags is None:
            input_tags = self._infer_tags(input_signal)
        
        # 匹配: 输入标签与单元标签的交集
        activated = []
        for uid, unit in self._registry.items():
            overlap = len(set(input_tags) & set(unit.tags))
            if overlap > 0:
                activated.append(uid)
                unit.activation_count += 1
        
        # 如果没有匹配的，激活负载最低的单元(兜底)
        if not activated:
            sorted_units = sorted(
                self._registry.values(), 
                key=lambda u: u.utilization
            )
            activated.append(sorted_units[0].unit_id)
            sorted_units[0].activation_count += 1
        
        # 路由冲突检测: 激活超过注册数60%说明维度太粗
        if len(activated) > len(self._registry) * self.DIMENSION_SPLIT_THRESHOLD:
            self._route_conflicts += 1
        self._total_dispatches += 1
        
        self._dispatch_count += 1
        return activated
    
    def check_split_needed(self) -> Optional[Tuple[str, Dict]]:
        """检查是否有单元需要拆分
        
        Returns:
            None 或 (unit_id, split_plan)
            split_plan: {
                'original_id': str,
                'child_a_id': str, 'child_a_tags': list,
                'child_b_id': str, 'child_b_tags': list,
                'overlap_tags': list,  # 重合标签
                'reason': str
            }
        """
        for uid, unit in self._registry.items():
            if unit.utilization >= self.split_threshold:
                plan = self._plan_split(unit)
                return uid, plan
        return None
    
    def _plan_split(self, unit: ManagedUnit) -> Dict:
        """规划拆分方案
        
        将一个满载单元拆成两个更专精的子单元:
        - 保留原有核心标签
        - 各自新增细化标签
        - 保留重合标签(过渡期不丢能力)
        """
        base_tags = unit.tags.copy()
        
        # 拆分策略: 保留所有原标签 + 各自新增细化方向标签
        # 核心原则: 拆分是细化，不是切割。子单元必须继承全部原标签
        #         + 各自新增方向标签，保证能力不丢失
        # 例: ['text','semantic'] → A:['text','semantic','text_A'], B:['text','semantic','text_B']
        import time
        ts = int(time.time())
        base = '_'.join(base_tags)
        child_a_tags = base_tags + [f"{base}_A"]
        child_b_tags = base_tags + [f"{base}_B"]
        # 重合: 全部原标签(完整继承)
        overlap_tags = base_tags.copy()
        
        plan = {
            'original_id': unit.unit_id,
            'original_tags': base_tags,
            'child_a_id': f"{unit.unit_id}_A_{ts}",
            'child_a_tags': child_a_tags,
            'child_b_id': f"{unit.unit_id}_B_{ts}",
            'child_b_tags': child_b_tags,
            'overlap_tags': overlap_tags,
            'capacity_each': unit.capacity,
            'reason': f'utilization={unit.utilization:.1%}'
        }
        
        self._split_history.append(plan)
        return plan
    
    def execute_split(self, plan: Dict):
        """执行拆分: 注销旧单元，注册两个子单元"""
        self.unregister(plan['original_id'])
        self.register(
            plan['child_a_id'], 
            plan['child_a_tags'],
            plan['capacity_each'],
            load=0  # 新生单元从0开始
        )
        self.register(
            plan['child_b_id'],
            plan['child_b_tags'], 
            plan['capacity_each'],
            load=0
        )
    
    def _infer_tags(self, signal: np.ndarray) -> List[str]:
        """从信号特征自动推断输入标签
        
        当前策略(文字优先):
        - 高稀疏(>80%零值) → text
        - 中稀疏(40-80%) → audio  
        - 低稀疏(<40%) → image
        """
        x = np.asarray(signal, dtype=np.float32).ravel()
        sparsity = float(np.count_nonzero(np.abs(x) < 0.01)) / max(len(x), 1)
        
        if sparsity > 0.8:
            return ['text', 'semantic']
        elif sparsity > 0.4:
            return ['audio', 'temporal']
        else:
            return ['image', 'spatial']
    
    def _rebuild_route_matrix(self):
        """重建路由矩阵(注册/注销后调用)"""
        n_units = len(self._registry)
        if n_units == 0:
            self._W_route = None
            return
        self._W_route = np.random.randn(
            n_units, self._tag_dim
        ).astype(np.float32) * 0.1
    
    @property
    def registry_info(self) -> Dict:
        """注册表摘要"""
        return {
            uid: {
                'tags': u.tags,
                'load': u.load,
                'capacity': u.capacity,
                'utilization': f'{u.utilization:.1%}',
                'activations': u.activation_count
            }
            for uid, u in self._registry.items()
        }
    
    @property
    def split_count(self) -> int:
        return len(self._split_history)
    
    def get_config(self) -> Dict:
        return {
            "type": "Manager",
            "version": "v2.0",
            "dimension": self.dimension,
            "num_neurons": self.num_neurons,
            "managed_units": len(self._registry),
            "split_count": self.split_count,
            "dimension_split_count": len(self._dimension_split_history),
            "dispatch_count": self._dispatch_count,
            "route_conflict_rate": self.route_conflict_rate
        }
    
    @property
    def route_conflict_rate(self) -> float:
        """路由冲突率: 越高说明维度越需要分裂"""
        if self._total_dispatches == 0:
            return 0.0
        return self._route_conflicts / self._total_dispatches
    
    def check_dimension_split_needed(self) -> Optional[Dict]:
        """检查管理微柱自身是否需要分裂
        
        触发条件: 路由冲突率超过阈值
        说明: 当前维度粒度太粗，总是同时激活太多不相关单元
        
        Returns:
            None 或 split_plan:
            {
                'original_dimension': str,
                'child_a_dimension': str, 'child_a_tags_subset': list,
                'child_b_dimension': str, 'child_b_tags_subset': list,
                'overlap_tags': list,
                'reason': str
            }
        """
        if self.route_conflict_rate < self.DIMENSION_SPLIT_THRESHOLD:
            return None
        if len(self._registry) < 2:
            return None  # 至少2个单元才有拆分意义
        
        # 按标签对注册单元做聚类，分两组
        all_tags = set()
        for unit in self._registry.values():
            all_tags.update(unit.tags)
        
        # 按标签频率排序，高频标签分到A组，低频分到B组
        tag_freq = {}
        for tag in all_tags:
            count = sum(1 for u in self._registry.values() if tag in u.tags)
            tag_freq[tag] = count
        sorted_tags = sorted(tag_freq.items(), key=lambda x: x[1], reverse=True)
        
        mid = len(sorted_tags) // 2
        tags_a = [t for t, _ in sorted_tags[:mid+1]]  # 高频组多1个保证重合
        tags_b = [t for t, _ in sorted_tags[mid:]]     # 低频组
        overlap = [t for t, _ in sorted_tags[max(0,mid-1):mid+2]]  # 中间重合
        
        import time
        ts = int(time.time())
        
        plan = {
            'original_dimension': self.dimension,
            'child_a_dimension': f"{self.dimension}_A",
            'child_a_tags_subset': tags_a,
            'child_b_dimension': f"{self.dimension}_B", 
            'child_b_tags_subset': tags_b,
            'overlap_tags': overlap,
            'reason': f'conflict_rate={self.route_conflict_rate:.1%}, units={len(self._registry)}'
        }
        
        self._dimension_split_history.append(plan)
        return plan