/**
 * Mixture of Experts (MoE) Router for Dynamic Agent Routing
 *
 * Features:
 * - 8 expert slots for specialized agent types
 * - Gating network for soft expert selection (top-k)
 * - Online weight updates via reward signals
 * - Load balancing with auxiliary loss
 * - Weight persistence to .swarm/moe-weights.json
 *
 * Architecture:
 * - Input: 384-dim task embedding (from ONNX)
 * - Hidden: 128-dim layer with ReLU
 * - Output: 8-dim softmax weights
 *
 * @module moe-router
 */
import { existsSync, mkdirSync, readFileSync, writeFileSync } from 'fs';
import { dirname } from 'path';
/**
 * Expert names in order (index corresponds to expert slot)
 */
export const EXPERT_NAMES = [
    'coder',
    'tester',
    'reviewer',
    'architect',
    'security',
    'performance',
    'researcher',
    'coordinator',
];
/**
 * Number of experts (fixed at 8)
 */
export const NUM_EXPERTS = 8;
/**
 * Input dimension (384 from ONNX MiniLM-L6-v2)
 */
export const INPUT_DIM = 384;
/**
 * Hidden layer dimension
 */
export const HIDDEN_DIM = 128;
/**
 * Default configuration
 */
const DEFAULT_CONFIG = {
    topK: 2,
    learningRate: 0.01,
    temperature: 1.0,
    loadBalanceCoef: 0.01,
    weightsPath: '.swarm/moe-weights.json',
    autoSaveInterval: 50,
    enableNoise: true,
    noiseStd: 0.1,
};
// ============================================================================
// Matrix Operations
// ============================================================================
/**
 * Initialize weights using Xavier/Glorot initialization
 */
function xavierInit(fanIn, fanOut) {
    const std = Math.sqrt(2.0 / (fanIn + fanOut));
    const weights = new Float32Array(fanIn * fanOut);
    for (let i = 0; i < weights.length; i++) {
        // Box-Muller transform for normal distribution
        const u1 = Math.random();
        const u2 = Math.random();
        const z = Math.sqrt(-2 * Math.log(u1 + 1e-8)) * Math.cos(2 * Math.PI * u2);
        weights[i] = z * std;
    }
    return weights;
}
/**
 * Matrix-vector multiplication: y = Wx
 * W is stored row-major: [rows * cols]
 */
function matmul(W, x, rows, cols, out) {
    for (let i = 0; i < rows; i++) {
        let sum = 0;
        const rowOffset = i * cols;
        // 4x loop unrolling for SIMD-friendly access
        let j = 0;
        for (; j + 3 < cols; j += 4) {
            sum +=
                W[rowOffset + j] * x[j] +
                    W[rowOffset + j + 1] * x[j + 1] +
                    W[rowOffset + j + 2] * x[j + 2] +
                    W[rowOffset + j + 3] * x[j + 3];
        }
        // Handle remainder
        for (; j < cols; j++) {
            sum += W[rowOffset + j] * x[j];
        }
        out[i] = sum;
    }
}
/**
 * Vector addition: y = x + b
 */
function addBias(x, b, out) {
    for (let i = 0; i < x.length; i++) {
        out[i] = x[i] + b[i];
    }
}
/**
 * ReLU activation: y = max(0, x)
 */
function relu(x, out) {
    for (let i = 0; i < x.length; i++) {
        out[i] = x[i] > 0 ? x[i] : 0;
    }
}
/**
 * Softmax with temperature: y_i = exp(x_i/T) / sum(exp(x_j/T))
 */
function softmax(x, temperature, out) {
    // Find max for numerical stability
    let maxVal = x[0];
    for (let i = 1; i < x.length; i++) {
        if (x[i] > maxVal)
            maxVal = x[i];
    }
    // Compute exp and sum
    let sum = 0;
    for (let i = 0; i < x.length; i++) {
        out[i] = Math.exp((x[i] - maxVal) / temperature);
        sum += out[i];
    }
    // Normalize
    const invSum = 1.0 / (sum + 1e-8);
    for (let i = 0; i < x.length; i++) {
        out[i] *= invSum;
    }
}
/**
 * Compute entropy of distribution: H = -sum(p * log(p))
 */
function entropy(p) {
    let h = 0;
    for (let i = 0; i < p.length; i++) {
        if (p[i] > 1e-8) {
            h -= p[i] * Math.log(p[i]);
        }
    }
    return h;
}
/**
 * Add Gaussian noise for exploration
 */
function addNoise(x, std, out) {
    for (let i = 0; i < x.length; i++) {
        const u1 = Math.random();
        const u2 = Math.random();
        const z = Math.sqrt(-2 * Math.log(u1 + 1e-8)) * Math.cos(2 * Math.PI * u2);
        out[i] = x[i] + z * std;
    }
}
// ============================================================================
// MoE Router Implementation
// ============================================================================
/**
 * Mixture of Experts Router
 *
 * Implements a two-layer gating network:
 * - Layer 1: Linear(384, 128) + ReLU
 * - Layer 2: Linear(128, 8) + Softmax
 *
 * Uses top-k expert selection with load balancing.
 */
export class MoERouter {
    config;
    // Network weights (pre-allocated Float32Arrays)
    W1; // [HIDDEN_DIM x INPUT_DIM]
    b1; // [HIDDEN_DIM]
    W2; // [NUM_EXPERTS x HIDDEN_DIM]
    b2; // [NUM_EXPERTS]
    // Intermediate buffers (pre-allocated, no GC pressure)
    hidden; // [HIDDEN_DIM]
    hiddenWithBias; // [HIDDEN_DIM]
    hiddenActivated; // [HIDDEN_DIM]
    logits; // [NUM_EXPERTS]
    logitsWithBias; // [NUM_EXPERTS]
    noisyLogits; // [NUM_EXPERTS]
    probs; // [NUM_EXPERTS]
    // Gradient buffers for online learning
    gradW2; // [NUM_EXPERTS x HIDDEN_DIM]
    gradb2; // [NUM_EXPERTS]
    gradW1; // [HIDDEN_DIM x INPUT_DIM]
    gradb1; // [HIDDEN_DIM]
    gradHidden; // [HIDDEN_DIM]
    // Statistics
    routingCounts; // [NUM_EXPERTS]
    totalRoutings = 0;
    updateCount = 0;
    avgReward = 0;
    // Cache for last input (for gradient computation)
    lastInput = null;
    lastHiddenActivated = null;
    lastProbs = null;
    lastSelectedExperts = [];
    constructor(config = {}) {
        this.config = { ...DEFAULT_CONFIG, ...config };
        // Initialize weights
        this.W1 = xavierInit(INPUT_DIM, HIDDEN_DIM);
        this.b1 = new Float32Array(HIDDEN_DIM);
        this.W2 = xavierInit(HIDDEN_DIM, NUM_EXPERTS);
        this.b2 = new Float32Array(NUM_EXPERTS);
        // Pre-allocate intermediate buffers
        this.hidden = new Float32Array(HIDDEN_DIM);
        this.hiddenWithBias = new Float32Array(HIDDEN_DIM);
        this.hiddenActivated = new Float32Array(HIDDEN_DIM);
        this.logits = new Float32Array(NUM_EXPERTS);
        this.logitsWithBias = new Float32Array(NUM_EXPERTS);
        this.noisyLogits = new Float32Array(NUM_EXPERTS);
        this.probs = new Float32Array(NUM_EXPERTS);
        // Pre-allocate gradient buffers
        this.gradW2 = new Float32Array(NUM_EXPERTS * HIDDEN_DIM);
        this.gradb2 = new Float32Array(NUM_EXPERTS);
        this.gradW1 = new Float32Array(HIDDEN_DIM * INPUT_DIM);
        this.gradb1 = new Float32Array(HIDDEN_DIM);
        this.gradHidden = new Float32Array(HIDDEN_DIM);
        // Statistics
        this.routingCounts = new Float32Array(NUM_EXPERTS);
    }
    /**
     * Initialize router, loading persisted weights if available
     */
    async initialize() {
        await this.loadWeights();
    }
    /**
     * Route task to top-k experts based on embedding
     *
     * @param taskEmbedding - 384-dim task embedding from ONNX
     * @returns Routing result with selected experts and weights
     */
    route(taskEmbedding) {
        // Convert to Float32Array if needed
        const input = taskEmbedding instanceof Float32Array
            ? taskEmbedding
            : new Float32Array(taskEmbedding);
        // Validate input dimension
        if (input.length !== INPUT_DIM) {
            throw new Error(`Expected embedding dimension ${INPUT_DIM}, got ${input.length}`);
        }
        // Forward pass through gating network
        // Layer 1: Linear + ReLU
        matmul(this.W1, input, HIDDEN_DIM, INPUT_DIM, this.hidden);
        addBias(this.hidden, this.b1, this.hiddenWithBias);
        relu(this.hiddenWithBias, this.hiddenActivated);
        // Layer 2: Linear
        matmul(this.W2, this.hiddenActivated, NUM_EXPERTS, HIDDEN_DIM, this.logits);
        addBias(this.logits, this.b2, this.logitsWithBias);
        // Add noise for exploration if enabled
        if (this.config.enableNoise) {
            addNoise(this.logitsWithBias, this.config.noiseStd, this.noisyLogits);
        }
        else {
            this.noisyLogits.set(this.logitsWithBias);
        }
        // Softmax to get probabilities
        softmax(this.noisyLogits, this.config.temperature, this.probs);
        // Select top-k experts
        const expertIndices = this.selectTopK(this.probs, this.config.topK);
        // Compute load balance loss
        const loadBalanceLoss = this.computeLoadBalanceLoss();
        // Compute entropy
        const routingEntropy = entropy(this.probs);
        // Update statistics
        for (const idx of expertIndices) {
            this.routingCounts[idx]++;
        }
        this.totalRoutings++;
        // Cache for gradient computation
        this.lastInput = new Float32Array(input);
        this.lastHiddenActivated = new Float32Array(this.hiddenActivated);
        this.lastProbs = new Float32Array(this.probs);
        this.lastSelectedExperts = expertIndices;
        // Build result
        const totalWeight = expertIndices.reduce((sum, idx) => sum + this.probs[idx], 0);
        const experts = expertIndices.map((idx) => ({
            name: EXPERT_NAMES[idx],
            index: idx,
            weight: this.probs[idx] / (totalWeight + 1e-8), // Normalize weights
            score: this.probs[idx],
        }));
        return {
            experts,
            allScores: Array.from(this.probs),
            loadBalanceLoss,
            entropy: routingEntropy,
        };
    }
    /**
     * Update expert weights based on reward signal
     *
     * Uses REINFORCE-style gradient update:
     * gradient = reward * d_log_prob / d_weights
     *
     * @param expert - Expert that received the reward
     * @param reward - Reward signal (-1 to 1, positive = good)
     */
    updateExpertWeights(expert, reward) {
        const expertIdx = typeof expert === 'number' ? expert : EXPERT_NAMES.indexOf(expert);
        if (expertIdx < 0 || expertIdx >= NUM_EXPERTS) {
            console.warn(`[MoE] Invalid expert: ${expert}`);
            return;
        }
        if (!this.lastInput || !this.lastHiddenActivated || !this.lastProbs) {
            console.warn('[MoE] No cached forward pass for gradient computation');
            return;
        }
        // Clamp reward to [-1, 1]
        const clampedReward = Math.max(-1, Math.min(1, reward));
        // Compute gradients using REINFORCE
        // For softmax: d_log_p_i / d_logit_j = delta_ij - p_j
        // gradient = reward * (1 - p_expert) for selected expert
        // gradient = reward * (-p_j) for other experts
        // Clear gradient buffers
        this.gradW2.fill(0);
        this.gradb2.fill(0);
        this.gradW1.fill(0);
        this.gradb1.fill(0);
        this.gradHidden.fill(0);
        // Gradient w.r.t. logits (before softmax)
        for (let i = 0; i < NUM_EXPERTS; i++) {
            if (i === expertIdx) {
                this.gradb2[i] = clampedReward * (1 - this.lastProbs[i]);
            }
            else {
                this.gradb2[i] = clampedReward * (-this.lastProbs[i]);
            }
        }
        // Gradient w.r.t. W2: outer product of gradb2 and hiddenActivated
        for (let i = 0; i < NUM_EXPERTS; i++) {
            const rowOffset = i * HIDDEN_DIM;
            for (let j = 0; j < HIDDEN_DIM; j++) {
                this.gradW2[rowOffset + j] = this.gradb2[i] * this.lastHiddenActivated[j];
            }
        }
        // Backprop through W2 to get gradient w.r.t. hidden
        for (let j = 0; j < HIDDEN_DIM; j++) {
            let sum = 0;
            for (let i = 0; i < NUM_EXPERTS; i++) {
                sum += this.gradb2[i] * this.W2[i * HIDDEN_DIM + j];
            }
            this.gradHidden[j] = sum;
        }
        // Backprop through ReLU: gradient is 0 where activation was 0
        for (let j = 0; j < HIDDEN_DIM; j++) {
            if (this.lastHiddenActivated[j] <= 0) {
                this.gradHidden[j] = 0;
            }
        }
        // Gradient w.r.t. b1
        this.gradb1.set(this.gradHidden);
        // Gradient w.r.t. W1: outer product of gradHidden and input
        for (let i = 0; i < HIDDEN_DIM; i++) {
            const rowOffset = i * INPUT_DIM;
            for (let j = 0; j < INPUT_DIM; j++) {
                this.gradW1[rowOffset + j] = this.gradHidden[i] * this.lastInput[j];
            }
        }
        // Apply gradients with learning rate
        const lr = this.config.learningRate;
        for (let i = 0; i < this.W2.length; i++) {
            this.W2[i] += lr * this.gradW2[i];
        }
        for (let i = 0; i < this.b2.length; i++) {
            this.b2[i] += lr * this.gradb2[i];
        }
        for (let i = 0; i < this.W1.length; i++) {
            this.W1[i] += lr * this.gradW1[i];
        }
        for (let i = 0; i < this.b1.length; i++) {
            this.b1[i] += lr * this.gradb1[i];
        }
        // Update statistics
        this.updateCount++;
        this.avgReward = (this.avgReward * (this.updateCount - 1) + clampedReward) / this.updateCount;
        // Auto-save
        if (this.config.autoSaveInterval > 0 &&
            this.updateCount % this.config.autoSaveInterval === 0) {
            this.saveWeights().catch(() => { }); // Fire and forget
        }
    }
    /**
     * Get load balance statistics across all experts
     */
    getLoadBalance() {
        const counts = {};
        const utilization = {};
        const total = this.totalRoutings || 1;
        const idealUtilization = 1 / NUM_EXPERTS;
        for (let i = 0; i < NUM_EXPERTS; i++) {
            const name = EXPERT_NAMES[i];
            counts[name] = this.routingCounts[i];
            utilization[name] = this.routingCounts[i] / total;
        }
        // Compute Gini coefficient
        const gini = this.computeGiniCoefficient();
        // Compute coefficient of variation
        const mean = total / NUM_EXPERTS;
        let variance = 0;
        for (let i = 0; i < NUM_EXPERTS; i++) {
            variance += Math.pow(this.routingCounts[i] - mean, 2);
        }
        variance /= NUM_EXPERTS;
        const cv = Math.sqrt(variance) / (mean + 1e-8);
        return {
            utilization,
            totalRoutings: this.totalRoutings,
            routingCounts: counts,
            giniCoefficient: gini,
            coefficientOfVariation: cv,
        };
    }
    /**
     * Get router statistics
     */
    getStats() {
        return {
            totalRoutings: this.totalRoutings,
            updateCount: this.updateCount,
            avgReward: this.avgReward,
            topK: this.config.topK,
            temperature: this.config.temperature,
            learningRate: this.config.learningRate,
            giniCoefficient: this.computeGiniCoefficient(),
        };
    }
    /**
     * Reset all statistics and routing counts
     */
    resetStats() {
        this.routingCounts.fill(0);
        this.totalRoutings = 0;
        this.updateCount = 0;
        this.avgReward = 0;
    }
    /**
     * Load weights from persistence file
     */
    async loadWeights(path) {
        const weightsPath = path || this.config.weightsPath;
        try {
            if (!existsSync(weightsPath)) {
                return false;
            }
            const data = readFileSync(weightsPath, 'utf-8');
            const model = JSON.parse(data);
            // Validate version
            if (!model.version || !model.version.startsWith('1.')) {
                console.warn(`[MoE] Incompatible model version: ${model.version}`);
                return false;
            }
            // Load weights
            this.W1 = new Float32Array(model.weights.W1.flat());
            this.b1 = new Float32Array(model.weights.b1);
            this.W2 = new Float32Array(model.weights.W2.flat());
            this.b2 = new Float32Array(model.weights.b2);
            // Load stats
            this.updateCount = model.stats.updateCount || 0;
            this.avgReward = model.stats.avgReward || 0;
            this.routingCounts = new Float32Array(model.stats.routingCounts || new Array(NUM_EXPERTS).fill(0));
            this.totalRoutings = this.routingCounts.reduce((a, b) => a + b, 0);
            return true;
        }
        catch (err) {
            console.warn(`[MoE] Failed to load weights: ${err}`);
            return false;
        }
    }
    /**
     * Save weights to persistence file
     */
    async saveWeights(path) {
        const weightsPath = path || this.config.weightsPath;
        try {
            // Ensure directory exists
            const dir = dirname(weightsPath);
            if (!existsSync(dir)) {
                mkdirSync(dir, { recursive: true });
            }
            // Convert Float32Arrays to nested arrays for JSON
            const W1_2d = [];
            for (let i = 0; i < HIDDEN_DIM; i++) {
                W1_2d.push(Array.from(this.W1.slice(i * INPUT_DIM, (i + 1) * INPUT_DIM)));
            }
            const W2_2d = [];
            for (let i = 0; i < NUM_EXPERTS; i++) {
                W2_2d.push(Array.from(this.W2.slice(i * HIDDEN_DIM, (i + 1) * HIDDEN_DIM)));
            }
            const model = {
                version: '1.0.0',
                config: {
                    topK: this.config.topK,
                    temperature: this.config.temperature,
                    learningRate: this.config.learningRate,
                    loadBalanceCoef: this.config.loadBalanceCoef,
                },
                weights: {
                    W1: W1_2d,
                    b1: Array.from(this.b1),
                    W2: W2_2d,
                    b2: Array.from(this.b2),
                },
                stats: {
                    updateCount: this.updateCount,
                    routingCounts: Array.from(this.routingCounts),
                    avgReward: this.avgReward,
                },
                metadata: {
                    savedAt: new Date().toISOString(),
                    expertNames: [...EXPERT_NAMES],
                },
            };
            writeFileSync(weightsPath, JSON.stringify(model, null, 2));
            return true;
        }
        catch (err) {
            console.warn(`[MoE] Failed to save weights: ${err}`);
            return false;
        }
    }
    /**
     * Reset weights to random initialization
     */
    resetWeights() {
        this.W1 = xavierInit(INPUT_DIM, HIDDEN_DIM);
        this.b1.fill(0);
        this.W2 = xavierInit(HIDDEN_DIM, NUM_EXPERTS);
        this.b2.fill(0);
        this.resetStats();
    }
    // ============================================================================
    // Private Methods
    // ============================================================================
    /**
     * Select top-k indices from probabilities
     */
    selectTopK(probs, k) {
        // Create index-value pairs and sort by value descending
        const pairs = [];
        for (let i = 0; i < probs.length; i++) {
            pairs.push([i, probs[i]]);
        }
        pairs.sort((a, b) => b[1] - a[1]);
        // Return top-k indices
        return pairs.slice(0, k).map((p) => p[0]);
    }
    /**
     * Compute load balance loss for regularization
     *
     * Uses auxiliary loss from Switch Transformer:
     * L_balance = N * sum(f_i * P_i)
     * where f_i = fraction of tokens routed to expert i
     *       P_i = average routing probability to expert i
     */
    computeLoadBalanceLoss() {
        if (this.totalRoutings === 0)
            return 0;
        let loss = 0;
        for (let i = 0; i < NUM_EXPERTS; i++) {
            const fraction = this.routingCounts[i] / this.totalRoutings;
            const avgProb = this.probs[i]; // Current routing prob
            loss += fraction * avgProb;
        }
        return NUM_EXPERTS * loss * this.config.loadBalanceCoef;
    }
    /**
     * Compute Gini coefficient for load distribution
     */
    computeGiniCoefficient() {
        if (this.totalRoutings === 0)
            return 0;
        // Sort counts
        const sorted = Array.from(this.routingCounts).sort((a, b) => a - b);
        const n = sorted.length;
        const mean = this.totalRoutings / n;
        // Compute Gini using the formula: G = (2 * sum(i * x_i) - (n+1) * sum(x_i)) / (n * sum(x_i))
        let weightedSum = 0;
        for (let i = 0; i < n; i++) {
            weightedSum += (i + 1) * sorted[i];
        }
        const gini = (2 * weightedSum - (n + 1) * this.totalRoutings) /
            (n * this.totalRoutings + 1e-8);
        return Math.max(0, gini);
    }
}
// ============================================================================
// Singleton Instance
// ============================================================================
let moeRouterInstance = null;
/**
 * Get singleton MoE router instance
 *
 * @param config - Optional configuration (only used on first call)
 * @returns MoE router instance
 */
export function getMoERouter(config) {
    if (!moeRouterInstance) {
        moeRouterInstance = new MoERouter(config);
        // Initialize in background (load weights)
        moeRouterInstance.initialize().catch((err) => {
            console.warn('[MoE] Failed to initialize router:', err);
        });
    }
    return moeRouterInstance;
}
/**
 * Reset singleton instance (for testing)
 */
export function resetMoERouter() {
    moeRouterInstance = null;
}
/**
 * Factory function to create new router
 */
export function createMoERouter(config) {
    return new MoERouter(config);
}
//# sourceMappingURL=moe-router.js.map