/**
 * Agentic QE v3 - Neural Topology Optimizer Tests
 * ADR-034: Neural Topology Optimizer
 *
 * Tests for the NeuralTopologyOptimizer implementation.
 */

import { describe, it, expect, beforeEach } from 'vitest';
import {
  NeuralTopologyOptimizer,
  createNeuralTopologyOptimizer,
} from '../../../src/neural-optimizer/topology-optimizer';
import {
  MutableSwarmTopology,
  createTopology,
  createAgent,
  buildMeshTopology,
  buildRingTopology,
} from '../../../src/neural-optimizer/swarm-topology';
import {
  DEFAULT_OPTIMIZER_CONFIG,
  actionToIndex,
  indexToActionType,
  ACTION_TYPES,
} from '../../../src/neural-optimizer/types';
import type { TopologyAgent } from '../../../src/neural-optimizer/types';

function createTestTopology(numAgents: number = 5): MutableSwarmTopology {
  const topology = createTopology('custom');

  for (let i = 0; i < numAgents; i++) {
    topology.addAgent(
      createAgent(`agent-${i}`, 'worker', {
        metrics: {
          tasksCompleted: Math.floor(Math.random() * 100),
          avgTaskDurationMs: Math.random() * 1000,
          successRate: 0.8 + Math.random() * 0.2,
          currentLoad: Math.random() * 0.6,
        },
      })
    );
  }

  // Add some initial connections
  for (let i = 0; i < numAgents - 1; i++) {
    topology.addConnection(`agent-${i}`, `agent-${i + 1}`, Math.random());
  }

  return topology;
}

describe('NeuralTopologyOptimizer', () => {
  let topology: MutableSwarmTopology;
  let optimizer: NeuralTopologyOptimizer;

  beforeEach(() => {
    topology = createTestTopology(5);
    optimizer = new NeuralTopologyOptimizer(topology);
  });

  describe('initialization', () => {
    it('should create optimizer with default config', () => {
      const stats = optimizer.getStats();

      expect(stats.totalSteps).toBe(0);
      expect(stats.episodes).toBe(0);
      expect(stats.currentEpsilon).toBe(DEFAULT_OPTIMIZER_CONFIG.epsilon);
    });

    it('should accept custom config', () => {
      const customOptimizer = new NeuralTopologyOptimizer(topology, {
        epsilon: 0.5,
        learningRate: 0.01,
        hiddenSize: 128,
      });

      const stats = customOptimizer.getStats();
      expect(stats.currentEpsilon).toBe(0.5);
    });

    it('should handle empty topology', () => {
      const emptyTopology = createTopology();
      const emptyOptimizer = new NeuralTopologyOptimizer(emptyTopology);

      // Should not throw
      const result = emptyOptimizer.optimizeStep();
      expect(result).toBeDefined();
    });
  });

  describe('optimizeStep', () => {
    it('should return optimization result', () => {
      const result = optimizer.optimizeStep();

      expect(result.action).toBeDefined();
      expect(typeof result.reward).toBe('number');
      expect(typeof result.newMinCut).toBe('number');
      expect(typeof result.communicationEfficiency).toBe('number');
      expect(typeof result.loadBalance).toBe('number');
      expect(typeof result.tdError).toBe('number');
      expect(typeof result.epsilon).toBe('number');
    });

    it('should increment total steps', () => {
      const statsBefore = optimizer.getStats();

      optimizer.optimizeStep();

      const statsAfter = optimizer.getStats();
      expect(statsAfter.totalSteps).toBe(statsBefore.totalSteps + 1);
    });

    it('should decay epsilon over time', () => {
      const initialEpsilon = optimizer.getStats().currentEpsilon;

      for (let i = 0; i < 10; i++) {
        optimizer.optimizeStep();
      }

      const finalEpsilon = optimizer.getStats().currentEpsilon;
      expect(finalEpsilon).toBeLessThan(initialEpsilon);
    });

    it('should not decay epsilon below minimum', () => {
      const minEpsilon = 0.01;
      const fastDecayOptimizer = new NeuralTopologyOptimizer(topology, {
        epsilon: 0.1,
        epsilonDecay: 0.5,
        minEpsilon,
      });

      for (let i = 0; i < 100; i++) {
        fastDecayOptimizer.optimizeStep();
      }

      expect(fastDecayOptimizer.getStats().currentEpsilon).toBeGreaterThanOrEqual(
        minEpsilon
      );
    });

    it('should update action counts', () => {
      const statsBefore = optimizer.getStats();
      const totalActionsBefore = Object.values(statsBefore.actionCounts).reduce(
        (sum, count) => sum + count,
        0
      );

      optimizer.optimizeStep();

      const statsAfter = optimizer.getStats();
      const totalActionsAfter = Object.values(statsAfter.actionCounts).reduce(
        (sum, count) => sum + count,
        0
      );

      expect(totalActionsAfter).toBe(totalActionsBefore + 1);
    });

    it('should track rewards', () => {
      optimizer.optimizeStep();

      const stats = optimizer.getStats();
      expect(stats.rewardHistory).toHaveLength(1);
    });

    it('should track min-cut history', () => {
      optimizer.optimizeStep();

      const stats = optimizer.getStats();
      expect(stats.minCutHistory).toHaveLength(1);
    });
  });

  describe('optimize (multiple steps)', () => {
    it('should run multiple optimization steps', () => {
      const results = optimizer.optimize(10);

      expect(results).toHaveLength(10);
      expect(optimizer.getStats().totalSteps).toBe(10);
    });

    it('should increment episode count', () => {
      optimizer.optimize(10);

      expect(optimizer.getStats().episodes).toBe(1);

      optimizer.optimize(10);

      expect(optimizer.getStats().episodes).toBe(2);
    });

    it('should accumulate cumulative reward', () => {
      const results = optimizer.optimize(10);
      const totalReward = results.reduce((sum, r) => sum + r.reward, 0);

      const stats = optimizer.getStats();
      expect(Math.abs(stats.cumulativeReward - totalReward)).toBeLessThan(0.001);
    });
  });

  describe('provideFeedback', () => {
    it('should accept external feedback', () => {
      optimizer.optimizeStep();

      // Should not throw
      optimizer.provideFeedback(1.0);
    });

    it('should update learning based on feedback', () => {
      // Run some steps first
      for (let i = 0; i < 5; i++) {
        optimizer.optimizeStep();
      }

      // Provide strong positive feedback
      optimizer.provideFeedback(1.0);

      // Stats should be updated
      const stats = optimizer.getStats();
      expect(stats.totalSteps).toBeGreaterThan(0);
    });
  });

  describe('getSkipRegions', () => {
    it('should return low connectivity agents', () => {
      // Create topology with isolated agent
      const sparseTopology = createTopology();
      sparseTopology.addAgent(createAgent('connected-1', 'worker'));
      sparseTopology.addAgent(createAgent('connected-2', 'worker'));
      sparseTopology.addAgent(createAgent('isolated', 'worker'));
      sparseTopology.addConnection('connected-1', 'connected-2');

      const sparseOptimizer = new NeuralTopologyOptimizer(sparseTopology);
      const skipRegions = sparseOptimizer.getSkipRegions();

      // Isolated agent should be in skip regions (degree < 2)
      expect(skipRegions).toContain('isolated');
    });

    it('should return empty array for well-connected topology', () => {
      const meshTopology = buildMeshTopology([
        createAgent('agent-1', 'worker'),
        createAgent('agent-2', 'worker'),
        createAgent('agent-3', 'worker'),
      ]);

      const meshOptimizer = new NeuralTopologyOptimizer(meshTopology);
      const skipRegions = meshOptimizer.getSkipRegions();

      expect(skipRegions).toHaveLength(0);
    });
  });

  describe('getStats', () => {
    it('should return valid statistics', () => {
      optimizer.optimize(10);

      const stats = optimizer.getStats();

      expect(stats.totalSteps).toBe(10);
      expect(stats.episodes).toBe(1);
      expect(typeof stats.cumulativeReward).toBe('number');
      expect(typeof stats.avgReward).toBe('number');
      expect(typeof stats.avgTdError).toBe('number');
      expect(stats.minCutHistory).toHaveLength(10);
      expect(stats.rewardHistory).toHaveLength(10);
    });

    it('should calculate average reward correctly', () => {
      const results = optimizer.optimize(10);
      const expectedAvg =
        results.reduce((sum, r) => sum + r.reward, 0) / results.length;

      const stats = optimizer.getStats();
      expect(Math.abs(stats.avgReward - expectedAvg)).toBeLessThan(0.001);
    });

    it('should track all action types', () => {
      optimizer.optimize(100);

      const stats = optimizer.getStats();

      for (const actionType of ACTION_TYPES) {
        expect(actionType in stats.actionCounts).toBe(true);
      }
    });
  });

  describe('reset', () => {
    it('should reset optimizer state', () => {
      optimizer.optimize(10);
      optimizer.reset();

      const stats = optimizer.getStats();
      expect(stats.currentEpsilon).toBe(DEFAULT_OPTIMIZER_CONFIG.epsilon);
    });

    it('should preserve learned weights', () => {
      optimizer.optimize(50);
      const modelBefore = optimizer.exportModel();

      optimizer.reset();

      const modelAfter = optimizer.exportModel();
      expect(modelAfter.valueNetwork.wHidden[0][0]).toBe(
        modelBefore.valueNetwork.wHidden[0][0]
      );
    });
  });

  describe('hardReset', () => {
    it('should clear all learning', () => {
      optimizer.optimize(10);
      optimizer.hardReset();

      const stats = optimizer.getStats();
      expect(stats.totalSteps).toBe(0);
      expect(stats.episodes).toBe(0);
      expect(stats.cumulativeReward).toBe(0);
      expect(stats.minCutHistory).toHaveLength(0);
    });

    it('should reinitialize networks', () => {
      optimizer.optimize(50);
      const modelBefore = optimizer.exportModel();

      optimizer.hardReset();

      const modelAfter = optimizer.exportModel();
      // Weights should be different after reinitialization
      expect(modelAfter.valueNetwork.wHidden[0][0]).not.toBe(
        modelBefore.valueNetwork.wHidden[0][0]
      );
    });
  });

  describe('exportModel', () => {
    it('should export valid model', () => {
      optimizer.optimize(10);

      const model = optimizer.exportModel();

      expect(model.type).toBe('neural-topology-optimizer');
      expect(model.version).toBeDefined();
      expect(model.config).toBeDefined();
      expect(model.valueNetwork).toBeDefined();
      expect(model.targetNetwork).toBeDefined();
      expect(model.stats).toBeDefined();
      expect(model.exportedAt).toBeDefined();
    });

    it('should include training stats', () => {
      optimizer.optimize(10);

      const model = optimizer.exportModel();

      expect(model.stats.totalSteps).toBe(10);
      expect(model.stats.episodes).toBe(1);
    });

    it('should include network weights', () => {
      optimizer.optimize(10);

      const model = optimizer.exportModel();

      expect(model.valueNetwork.wHidden).toBeDefined();
      expect(model.valueNetwork.bHidden).toBeDefined();
      expect(model.valueNetwork.wOutput).toBeDefined();
      expect(typeof model.valueNetwork.bOutput).toBe('number');
    });
  });

  describe('importModel', () => {
    it('should import exported model', () => {
      optimizer.optimize(50);
      const model = optimizer.exportModel();

      const newTopology = createTestTopology(5);
      const newOptimizer = new NeuralTopologyOptimizer(newTopology);
      newOptimizer.importModel(model);

      const newStats = newOptimizer.getStats();
      expect(newStats.totalSteps).toBe(50);
      expect(newStats.currentEpsilon).toBe(model.stats.currentEpsilon);
    });

    it('should throw on invalid model type', () => {
      const invalidModel = {
        type: 'invalid-type',
        version: '1.0.0',
        config: DEFAULT_OPTIMIZER_CONFIG,
        valueNetwork: { wHidden: [[0]], bHidden: [0], wOutput: [0], bOutput: 0 },
        stats: optimizer.getStats(),
        exportedAt: new Date().toISOString(),
      };

      expect(() => optimizer.importModel(invalidModel as any)).toThrow();
    });

    it('should reproduce same behavior after import', () => {
      // Train original
      optimizer.optimize(100);

      // Export and import
      const model = optimizer.exportModel();
      const newTopology = createTestTopology(5);
      const newOptimizer = new NeuralTopologyOptimizer(newTopology);
      newOptimizer.importModel(model);

      // Both should have same epsilon
      expect(newOptimizer.getStats().currentEpsilon).toBe(
        optimizer.getStats().currentEpsilon
      );
    });
  });

  describe('topology modifications', () => {
    it('should add connections during optimization', () => {
      const sparseTopology = createTopology();
      sparseTopology.addAgent(createAgent('agent-1', 'worker'));
      sparseTopology.addAgent(createAgent('agent-2', 'worker'));
      sparseTopology.addAgent(createAgent('agent-3', 'worker'));
      // No initial connections

      const sparseOptimizer = new NeuralTopologyOptimizer(sparseTopology, {
        epsilon: 1.0, // Force exploration
      });

      const initialConnections = sparseTopology.connections.length;

      // Run many steps to ensure some add_connection actions
      for (let i = 0; i < 100; i++) {
        sparseOptimizer.optimizeStep();
      }

      const stats = sparseOptimizer.getStats();
      expect(stats.actionCounts['add_connection']).toBeGreaterThan(0);
    });

    it('should remove connections during optimization', () => {
      const denseTopology = buildMeshTopology([
        createAgent('agent-1', 'worker'),
        createAgent('agent-2', 'worker'),
        createAgent('agent-3', 'worker'),
        createAgent('agent-4', 'worker'),
      ]);

      const denseOptimizer = new NeuralTopologyOptimizer(denseTopology, {
        epsilon: 1.0, // Force exploration
      });

      for (let i = 0; i < 100; i++) {
        denseOptimizer.optimizeStep();
      }

      const stats = denseOptimizer.getStats();
      expect(stats.actionCounts['remove_connection']).toBeGreaterThan(0);
    });

    it('should modify connection weights', () => {
      optimizer.optimize(100);

      const stats = optimizer.getStats();
      const weightModifications =
        stats.actionCounts['strengthen_connection'] +
        stats.actionCounts['weaken_connection'];

      expect(weightModifications).toBeGreaterThan(0);
    });
  });

  describe('learning behavior', () => {
    it('should learn from experience replay', () => {
      const learningOptimizer = new NeuralTopologyOptimizer(topology, {
        minExperiencesForTraining: 10,
        batchSize: 8,
      });

      // Generate enough experiences
      for (let i = 0; i < 20; i++) {
        learningOptimizer.optimizeStep();
      }

      // Should have trained from replay
      const stats = learningOptimizer.getStats();
      expect(stats.totalSteps).toBe(20);
    });

    it('should converge value estimates over time', () => {
      const learningOptimizer = new NeuralTopologyOptimizer(topology, {
        learningRate: 0.01,
        minExperiencesForTraining: 50,
      });

      // Track TD errors over time
      const earlyTdErrors: number[] = [];
      const lateTdErrors: number[] = [];

      for (let i = 0; i < 200; i++) {
        const result = learningOptimizer.optimizeStep();
        if (i < 50) {
          earlyTdErrors.push(Math.abs(result.tdError));
        } else if (i >= 150) {
          lateTdErrors.push(Math.abs(result.tdError));
        }
      }

      const earlyAvg = earlyTdErrors.reduce((a, b) => a + b, 0) / earlyTdErrors.length;
      const lateAvg = lateTdErrors.reduce((a, b) => a + b, 0) / lateTdErrors.length;

      // TD errors should generally decrease (value estimates improve)
      // Note: This test is probabilistic and may occasionally fail
      expect(lateAvg).toBeLessThanOrEqual(earlyAvg * 2); // Allow some variance
    });
  });

  describe('multi-objective reward', () => {
    it('should consider efficiency in reward', () => {
      const efficiencyOptimizer = new NeuralTopologyOptimizer(topology, {
        efficiencyWeight: 1.0,
        loadBalanceWeight: 0,
        latencyWeight: 0,
      });

      const result = efficiencyOptimizer.optimizeStep();

      expect(typeof result.communicationEfficiency).toBe('number');
      expect(result.communicationEfficiency).toBeGreaterThanOrEqual(0);
      expect(result.communicationEfficiency).toBeLessThanOrEqual(1);
    });

    it('should consider load balance in reward', () => {
      const loadOptimizer = new NeuralTopologyOptimizer(topology, {
        efficiencyWeight: 0,
        loadBalanceWeight: 1.0,
        latencyWeight: 0,
      });

      const result = loadOptimizer.optimizeStep();

      expect(typeof result.loadBalance).toBe('number');
    });

    it('should bound reward to [-1, 1]', () => {
      for (let i = 0; i < 50; i++) {
        const result = optimizer.optimizeStep();

        expect(result.reward).toBeGreaterThanOrEqual(-1);
        expect(result.reward).toBeLessThanOrEqual(1);
      }
    });
  });

  describe('factory function', () => {
    it('createNeuralTopologyOptimizer should create optimizer', () => {
      const factoryOptimizer = createNeuralTopologyOptimizer(topology, {
        epsilon: 0.2,
      });

      const stats = factoryOptimizer.getStats();
      expect(stats.currentEpsilon).toBe(0.2);
    });
  });
});

describe('action helpers', () => {
  describe('actionToIndex', () => {
    it('should map actions to indices', () => {
      expect(actionToIndex({ type: 'add_connection', from: 'a', to: 'b' })).toBe(0);
      expect(actionToIndex({ type: 'remove_connection', from: 'a', to: 'b' })).toBe(1);
      expect(
        actionToIndex({ type: 'strengthen_connection', from: 'a', to: 'b', delta: 0.1 })
      ).toBe(2);
      expect(
        actionToIndex({ type: 'weaken_connection', from: 'a', to: 'b', delta: 0.1 })
      ).toBe(3);
      expect(actionToIndex({ type: 'no_op' })).toBe(4);
    });
  });

  describe('indexToActionType', () => {
    it('should map indices to action types', () => {
      expect(indexToActionType(0)).toBe('add_connection');
      expect(indexToActionType(1)).toBe('remove_connection');
      expect(indexToActionType(2)).toBe('strengthen_connection');
      expect(indexToActionType(3)).toBe('weaken_connection');
      expect(indexToActionType(4)).toBe('no_op');
    });

    it('should wrap around for large indices', () => {
      expect(indexToActionType(5)).toBe('add_connection');
      expect(indexToActionType(6)).toBe('remove_connection');
    });
  });
});