"""虫群训练器 — 统一训练入口"""
import numpy as np
import os
import json
import time
from typing import Dict, List, Optional, Tuple


class SwarmTrainer:
    """统一训练器 — 管理脑区权重训练"""
    
    def __init__(self, brain, learning_rate: float = 0.01):
        self.brain = brain
        self.lr = learning_rate
        self._step = 0
        self._history = []  # 训练历史
    
    def train_qa(self, question: str, answer: str) -> Dict:
        """训练QA对 — 写入记忆+微调权重"""
        t0 = time.time()
        result = {
            'question': question,
            'answer': answer,
            'areas_updated': [],
            'weight_changes': {},
        }
        
        # 1. 写入语义记忆
        if hasattr(self.brain, 'memory') and self.brain.memory:
            self.brain.memory.store_qa(question, answer)
            result['areas_updated'].append('memory')
        
        # 2. 写入QA缓存
        if hasattr(self.brain, 'chat_engine') and self.brain.chat_engine:
            self.brain.chat_engine._add_to_cache(question, answer)
            result['areas_updated'].append('chat_cache')
        
        # 3. 微调相关脑区权重
        try:
            q_vec = self.brain.embed_text(question)
            a_vec = self.brain.embed_text(answer)
            
            # 赫布学习: 同时激活的连接加强
            for area_name, area in self.brain.areas.items():
                if hasattr(area, 'hebbian_update'):
                    change = area.hebbian_update(q_vec, a_vec, self.lr)
                    if change > 1e-6:
                        result['weight_changes'][area_name] = float(change)
                        if area_name not in result['areas_updated']:
                            result['areas_updated'].append(area_name)
        except Exception as e:
            result['hebbian_error'] = str(e)
        
        self._step += 1
        result['ms'] = int((time.time() - t0) * 1000)
        result['step'] = self._step
        
        self._history.append(result)
        return result
    
    def train_batch(self, qa_pairs: List[Tuple[str, str]]) -> Dict:
        """批量训练QA对"""
        results = []
        for q, a in qa_pairs:
            r = self.train_qa(q, a)
            results.append(r)
        
        total_changes = sum(
            len(r.get('weight_changes', {})) for r in results
        )
        
        return {
            'total': len(qa_pairs),
            'areas_updated': len(set(
                a for r in results for a in r.get('areas_updated', [])
            )),
            'weight_changes': total_changes,
            'avg_ms': sum(r['ms'] for r in results) // max(len(results), 1),
        }
    
    def get_status(self) -> Dict:
        """训练状态"""
        return {
            'step': self._step,
            'learning_rate': self.lr,
            'history_size': len(self._history),
        }