#!/usr/bin/env python3
"""
维度适配器 — 打通Swarm(32维)和Meta Model(75维)的信号通路

双向转换:
- up(32→75): 自编码器扩展 + 残差连接
- down(75→32): 自编码器压缩 + 残差连接

训练: 可通过train()方法学习双向映射
"""

import numpy as np
import os
import pickle
from typing import Dict, Tuple


class DimensionAdapter:
    """
    维度适配器 — 32维↔75维双向转换
    
    设计:
    - 双层自编码器(32→64→75 / 75→64→32)
    - 残差连接防止信号丢失
    - 可训练(通过train()方法)
    """
    
    def __init__(self, swarm_dim: int = 32, meta_dim: int = 75):
        self.swarm_dim = swarm_dim
        self.meta_dim = meta_dim
        
        # 中间维度
        self.mid_dim = 64
        
        # 上投影: swarm_dim → mid_dim → meta_dim
        self.up_W1 = np.random.randn(self.mid_dim, swarm_dim).astype(np.float32) * 0.1
        self.up_b1 = np.zeros(self.mid_dim, dtype=np.float32)
        self.up_W2 = np.random.randn(meta_dim, self.mid_dim).astype(np.float32) * 0.1
        self.up_b2 = np.zeros(meta_dim, dtype=np.float32)
        
        # 下投影: meta_dim → mid_dim → swarm_dim
        self.down_W1 = np.random.randn(self.mid_dim, meta_dim).astype(np.float32) * 0.1
        self.down_b1 = np.zeros(self.mid_dim, dtype=np.float32)
        self.down_W2 = np.random.randn(swarm_dim, self.mid_dim).astype(np.float32) * 0.1
        self.down_b2 = np.zeros(swarm_dim, dtype=np.float32)
        
        # 训练状态
        self._trained = False
        self._train_loss = None
    
    def up(self, x: np.ndarray) -> np.ndarray:
        """32维 → 75维 (Swarm → Meta Model)"""
        x = np.asarray(x, dtype=np.float32).ravel()
        if len(x) < self.swarm_dim:
            x = np.pad(x, (0, self.swarm_dim - len(x)))
        elif len(x) > self.swarm_dim:
            x = x[:self.swarm_dim]
        
        # 双层变换 + ReLU
        h = np.maximum(0, self.up_W1 @ x + self.up_b1)
        out = self.up_W2 @ h + self.up_b2
        
        # 残差: 原始信号投影到75维后叠加
        residual = np.zeros(self.meta_dim, dtype=np.float32)
        residual[:self.swarm_dim] = x * 0.2  # 前32维保留20%原始信号
        out = out + residual
        
        # 归一化
        norm = np.linalg.norm(out)
        if norm > 1e-8:
            out = out / norm * np.sqrt(self.meta_dim)
        
        return out
    
    def down(self, x: np.ndarray) -> np.ndarray:
        """75维 → 32维 (Meta Model → Swarm)"""
        x = np.asarray(x, dtype=np.float32).ravel()
        if len(x) < self.meta_dim:
            x = np.pad(x, (0, self.meta_dim - len(x)))
        elif len(x) > self.meta_dim:
            x = x[:self.meta_dim]
        
        # 双层变换 + ReLU
        h = np.maximum(0, self.down_W1 @ x + self.down_b1)
        out = self.down_W2 @ h + self.down_b2
        
        # 残差: 从75维中取前32维保留
        residual = x[:self.swarm_dim] * 0.2
        out = out + residual
        
        # 归一化
        norm = np.linalg.norm(out)
        if norm > 1e-8:
            out = out / norm * np.sqrt(self.swarm_dim)
        
        return out
    
    def train(self, n_samples: int = 200, epochs: int = 3, lr: float = 0.01):
        """自编码器训练 — 学习双向映射"""
        # 生成训练数据(随机向量模拟embedding)
        np.random.seed(42)
        samples = np.random.randn(n_samples, self.swarm_dim).astype(np.float32) * 0.5
        
        for epoch in range(epochs):
            total_loss = 0.0
            for x in samples:
                # 前向: 32→75→32
                h_up = np.maximum(0, self.up_W1 @ x + self.up_b1)
                mid = self.up_W2 @ h_up + self.up_b2
                
                h_down = np.maximum(0, self.down_W1 @ mid + self.down_b1)
                recon = self.down_W2 @ h_down + self.down_b2
                
                # 重建误差
                loss = float(np.mean((recon - x) ** 2))
                total_loss += loss
                
                # 简单梯度下降(数值梯度近似)
                grad_scale = lr * 2 * (recon - x) / n_samples
                self.down_W2 -= np.outer(grad_scale, h_down)
                self.down_b2 -= grad_scale
            
            avg_loss = total_loss / n_samples
            self._train_loss = avg_loss
        
        self._trained = True
    
    def info(self) -> Dict:
        """获取适配器信息"""
        up_params = self.up_W1.size + self.up_b1.size + self.up_W2.size + self.up_b2.size
        down_params = self.down_W1.size + self.down_b1.size + self.down_W2.size + self.down_b2.size
        return {
            'type': 'DimensionAdapter',
            'swarm_dim': self.swarm_dim,
            'meta_dim': self.meta_dim,
            'params': up_params + down_params,
            'up_count': 2,
            'down_count': 2,
            'trained': self._trained,
            'train_loss': self._train_loss,
        }
    
    def save(self, path: str):
        """保存适配器"""
        data = {
            'swarm_dim': self.swarm_dim,
            'meta_dim': self.meta_dim,
            'up_W1': self.up_W1, 'up_b1': self.up_b1,
            'up_W2': self.up_W2, 'up_b2': self.up_b2,
            'down_W1': self.down_W1, 'down_b1': self.down_b1,
            'down_W2': self.down_W2, 'down_b2': self.down_b2,
            'trained': self._trained,
            'train_loss': self._train_loss,
        }
        with open(path, 'wb') as f:
            pickle.dump(data, f)
    
    @classmethod
    def load(cls, path: str) -> 'DimensionAdapter':
        """加载适配器"""
        with open(path, 'rb') as f:
            data = pickle.load(f)
        adapter = cls(swarm_dim=data['swarm_dim'], meta_dim=data['meta_dim'])
        adapter.up_W1 = data['up_W1']
        adapter.up_b1 = data['up_b1']
        adapter.up_W2 = data['up_W2']
        adapter.up_b2 = data['up_b2']
        adapter.down_W1 = data['down_W1']
        adapter.down_b1 = data['down_b1']
        adapter.down_W2 = data['down_W2']
        adapter.down_b2 = data['down_b2']
        adapter._trained = data.get('trained', False)
        adapter._train_loss = data.get('train_loss', None)
        return adapter