/**
 * Tests for NMF (Non-negative Matrix Factorization) module
 */

import { describe, it, expect, beforeEach } from 'vitest';
import { 
  NMFAlgorithm, 
  SpectralDictionary, 
  NMFUtils, 
  decomposeNMF,
  type NMFConfig,
  type NMFResult,
  type NMFComponent,
  type SpectralDictionary as SpectralDictionaryType
} from './nmf';

describe('NMFAlgorithm', () => {
  let nmf: NMFAlgorithm;
  let config: NMFConfig;

  beforeEach(() => {
    config = {
      rank: 3,
      maxIterations: 50,
      tolerance: 1e-4,
      sparsity: 0.1,
      smoothness: 0.1,
      useMultiplicative: true,
      useAlternating: false,
      useKullbackLeibler: false,
      useEuclidean: true,
      randomSeed: 42,
    };
    nmf = new NMFAlgorithm(config);
  });

  describe('constructor', () => {
    it('should create NMF with default config', () => {
      const defaultNMF = new NMFAlgorithm();
      const nmfConfig = defaultNMF.getConfig();
      
      expect(nmfConfig.rank).toBe(4);
      expect(nmfConfig.maxIterations).toBe(100);
      expect(nmfConfig.tolerance).toBe(1e-6);
      expect(nmfConfig.sparsity).toBe(0.1);
      expect(nmfConfig.smoothness).toBe(0.1);
      expect(nmfConfig.useMultiplicative).toBe(true);
      expect(nmfConfig.useAlternating).toBe(false);
      expect(nmfConfig.useKullbackLeibler).toBe(false);
      expect(nmfConfig.useEuclidean).toBe(true);
      expect(nmfConfig.randomSeed).toBe(42);
    });

    it('should accept custom config', () => {
      const customConfig: NMFConfig = {
        rank: 5,
        maxIterations: 200,
        tolerance: 1e-8,
        sparsity: 0.2,
        smoothness: 0.2,
        useMultiplicative: false,
        useAlternating: true,
        useKullbackLeibler: true,
        useEuclidean: false,
        randomSeed: 123,
      };
      
      const customNMF = new NMFAlgorithm(customConfig);
      const nmfConfig = customNMF.getConfig();
      
      expect(nmfConfig.rank).toBe(5);
      expect(nmfConfig.maxIterations).toBe(200);
      expect(nmfConfig.tolerance).toBe(1e-8);
      expect(nmfConfig.sparsity).toBe(0.2);
      expect(nmfConfig.smoothness).toBe(0.2);
      expect(nmfConfig.useMultiplicative).toBe(false);
      expect(nmfConfig.useAlternating).toBe(true);
      expect(nmfConfig.useKullbackLeibler).toBe(true);
      expect(nmfConfig.useEuclidean).toBe(false);
      expect(nmfConfig.randomSeed).toBe(123);
    });
  });

  describe('decompose', () => {
    it('should decompose simple matrix', () => {
      // Create simple test matrix
      const V = [
        new Float32Array([1, 2, 3, 4]),
        new Float32Array([2, 4, 6, 8]),
        new Float32Array([3, 6, 9, 12]),
      ];

      const result = nmf.decompose(V);
      
      expect(result.W).toHaveLength(3);
      expect(result.H).toHaveLength(3);
      expect(result.W[0]).toHaveLength(3);
      expect(result.H[0]).toHaveLength(4);
      expect(result.reconstruction).toHaveLength(3);
      expect(result.reconstruction[0]).toHaveLength(4);
      expect(result.error).toBeGreaterThanOrEqual(0);
      expect(result.iterations).toBeGreaterThan(0);
      expect(result.iterations).toBeLessThanOrEqual(50);
      expect(result.converged).toBeDefined();
      expect(result.components).toHaveLength(3);
    });

    it('should handle different matrix sizes', () => {
      const V = [
        new Float32Array([1, 2, 3]),
        new Float32Array([4, 5, 6]),
      ];

      const result = nmf.decompose(V);
      
      expect(result.W).toHaveLength(3);
      expect(result.H).toHaveLength(3);
      expect(result.W[0]).toHaveLength(2);
      expect(result.H[0]).toHaveLength(3);
    });

    it('should use provided initial matrices', () => {
      const V = [
        new Float32Array([1, 2, 3, 4]),
        new Float32Array([2, 4, 6, 8]),
      ];

      const initialW = [
        new Float32Array([0.5, 0.5]),
        new Float32Array([0.3, 0.7]),
        new Float32Array([0.8, 0.2]),
      ];

      const initialH = [
        new Float32Array([1, 2, 3, 4]),
        new Float32Array([2, 1, 4, 3]),
        new Float32Array([3, 4, 1, 2]),
      ];

      const result = nmf.decompose(V, initialW, initialH);
      
      expect(result.W).toHaveLength(3);
      expect(result.H).toHaveLength(3);
      expect(result.error).toBeGreaterThanOrEqual(0);
    });

    it('should converge within max iterations', () => {
      const V = [
        new Float32Array([1, 2, 3, 4]),
        new Float32Array([2, 4, 6, 8]),
        new Float32Array([3, 6, 9, 12]),
      ];

      const result = nmf.decompose(V);
      
      expect(result.iterations).toBeLessThanOrEqual(50);
      expect(result.error).toBeGreaterThanOrEqual(0);
    });

    it('should produce non-negative results', () => {
      const V = [
        new Float32Array([1, 2, 3, 4]),
        new Float32Array([2, 4, 6, 8]),
      ];

      const result = nmf.decompose(V);
      
      // Check W matrix
      for (const basis of result.W) {
        for (const value of basis) {
          expect(value).toBeGreaterThanOrEqual(0);
        }
      }
      
      // Check H matrix
      for (const coeffs of result.H) {
        for (const value of coeffs) {
          expect(value).toBeGreaterThanOrEqual(0);
        }
      }
    });
  });

  describe('configuration', () => {
    it('should update configuration', () => {
      const newConfig = {
        rank: 5,
        maxIterations: 200,
        tolerance: 1e-8,
      };

      nmf.updateConfig(newConfig);
      const nmfConfig = nmf.getConfig();
      
      expect(nmfConfig.rank).toBe(5);
      expect(nmfConfig.maxIterations).toBe(200);
      expect(nmfConfig.tolerance).toBe(1e-8);
    });

    it('should handle different algorithms', () => {
      const V = [
        new Float32Array([1, 2, 3, 4]),
        new Float32Array([2, 4, 6, 8]),
      ];

      // Test multiplicative updates
      const multiplicativeNMF = new NMFAlgorithm({
        useMultiplicative: true,
        useAlternating: false,
      });
      const result1 = multiplicativeNMF.decompose(V);
      expect(result1.error).toBeGreaterThanOrEqual(0);

      // Test alternating least squares
      const alternatingNMF = new NMFAlgorithm({
        useMultiplicative: false,
        useAlternating: true,
      });
      const result2 = alternatingNMF.decompose(V);
      expect(result2.error).toBeGreaterThanOrEqual(0);
    });
  });

  describe('component extraction', () => {
    it('should extract components with metadata', () => {
      const V = [
        new Float32Array([1, 2, 3, 4]),
        new Float32Array([2, 4, 6, 8]),
        new Float32Array([3, 6, 9, 12]),
      ];

      const result = nmf.decompose(V);
      
      expect(result.components).toHaveLength(3);
      
      for (const component of result.components) {
        expect(component.index).toBeGreaterThanOrEqual(0);
        expect(component.basis).toBeDefined();
        expect(component.coefficients).toBeDefined();
        expect(component.energy).toBeGreaterThanOrEqual(0);
        expect(component.spectralCentroid).toBeGreaterThanOrEqual(0);
        expect(component.spectralRolloff).toBeGreaterThanOrEqual(0);
        expect(component.fundamental).toBeGreaterThanOrEqual(0);
        expect(component.confidence).toBeGreaterThanOrEqual(0);
        expect(component.confidence).toBeLessThanOrEqual(1);
        expect(typeof component.isPitch).toBe('boolean');
      }
    });

    it('should identify pitch components', () => {
      const V = [
        new Float32Array([1, 2, 3, 4]),
        new Float32Array([2, 4, 6, 8]),
        new Float32Array([3, 6, 9, 12]),
      ];

      const result = nmf.decompose(V);
      
      // At least some components should be identified as pitches
      const pitchComponents = result.components.filter(comp => comp.isPitch);
      expect(pitchComponents.length).toBeGreaterThanOrEqual(0);
    });
  });
});

describe('SpectralDictionary', () => {
  describe('learnDictionary', () => {
    it('should learn dictionary from training data', () => {
      const trainingData = [
        [
          new Float32Array([1, 2, 3, 4]),
          new Float32Array([2, 4, 6, 8]),
        ],
        [
          new Float32Array([3, 6, 9, 12]),
          new Float32Array([4, 8, 12, 16]),
        ],
      ];

      const dictionary = SpectralDictionary.learnDictionary(trainingData, 5);
      
      expect(dictionary.atoms).toHaveLength(5);
      expect(dictionary.frequencies).toHaveLength(5);
      expect(dictionary.pitches).toHaveLength(5);
      expect(dictionary.weights).toHaveLength(5);
      expect(dictionary.size).toBe(5);
      
      for (const atom of dictionary.atoms) {
        expect(atom).toBeInstanceOf(Float32Array);
        expect(atom.length).toBeGreaterThan(0);
      }
      
      for (const freq of dictionary.frequencies) {
        expect(freq).toBeGreaterThan(0);
      }
      
      for (const pitch of dictionary.pitches) {
        expect(pitch).toBeGreaterThanOrEqual(0);
        expect(pitch).toBeLessThanOrEqual(127);
      }
      
      for (const weight of dictionary.weights) {
        expect(weight).toBeGreaterThanOrEqual(0);
      }
    });

    it('should handle empty training data', () => {
      const trainingData: Float32Array[][] = [];
      
      const dictionary = SpectralDictionary.learnDictionary(trainingData, 3);
      
      expect(dictionary.atoms).toHaveLength(0);
      expect(dictionary.size).toBe(0);
    });

    it('should use custom NMF config', () => {
      const trainingData = [
        [
          new Float32Array([1, 2, 3, 4]),
          new Float32Array([2, 4, 6, 8]),
        ],
      ];

      const customConfig = {
        rank: 3,
        maxIterations: 20,
        tolerance: 1e-4,
      };

      const dictionary = SpectralDictionary.learnDictionary(trainingData, 3, customConfig);
      
      expect(dictionary.size).toBe(3);
    });
  });

  describe('utility functions', () => {
    it('should flatten training data correctly', () => {
      const trainingData = [
        [
          new Float32Array([1, 2, 3]),
          new Float32Array([4, 5, 6]),
        ],
        [
          new Float32Array([7, 8, 9]),
        ],
      ];

      const flattened = SpectralDictionary.flattenTrainingData(trainingData);
      
      expect(flattened).toHaveLength(3);
      expect(flattened[0]).toEqual(new Float32Array([1, 2, 3]));
      expect(flattened[1]).toEqual(new Float32Array([4, 5, 6]));
      expect(flattened[2]).toEqual(new Float32Array([7, 8, 9]));
    });

    it('should estimate frequency from atom', () => {
      const atom = new Float32Array([0.1, 0.2, 0.8, 0.3, 0.1]);
      const frequency = SpectralDictionary.estimateFrequency(atom);
      
      expect(frequency).toBeGreaterThan(0);
    });

    it('should calculate atom weight', () => {
      const atom = new Float32Array([1, 2, 3, 4]);
      const weight = SpectralDictionary.calculateAtomWeight(atom);
      
      expect(weight).toBe(30); // 1² + 2² + 3² + 4² = 1 + 4 + 9 + 16 = 30
    });
  });
});

describe('NMFUtils', () => {
  let spectrogram: Float32Array[];
  let dictionary: SpectralDictionaryType;

  beforeEach(() => {
    spectrogram = [
      new Float32Array([1, 2, 3, 4]),
      new Float32Array([2, 4, 6, 8]),
      new Float32Array([3, 6, 9, 12]),
    ];

    dictionary = {
      atoms: [
        new Float32Array([0.5, 0.3, 0.2]),
        new Float32Array([0.2, 0.6, 0.2]),
        new Float32Array([0.1, 0.1, 0.8]),
      ],
      frequencies: [440, 880, 1320],
      pitches: [69, 81, 89],
      weights: [0.5, 0.3, 0.2],
      size: 3,
    };
  });

  describe('separateSources', () => {
    it('should separate sources using dictionary', () => {
      const result = NMFUtils.separateSources(spectrogram, dictionary);
      
      expect(result.W).toHaveLength(3);
      expect(result.H).toHaveLength(3);
      expect(result.error).toBeGreaterThanOrEqual(0);
      expect(result.components).toHaveLength(3);
    });

    it('should use custom NMF config', () => {
      const customConfig = {
        rank: 3,
        maxIterations: 20,
        tolerance: 1e-4,
      };

      const result = NMFUtils.separateSources(spectrogram, dictionary, customConfig);
      
      expect(result.W).toHaveLength(3);
      expect(result.H).toHaveLength(3);
    });
  });

  describe('extractPitchComponents', () => {
    it('should extract pitch components', () => {
      const result = {
        W: [new Float32Array([1, 2, 3])],
        H: [new Float32Array([1, 2, 3, 4])],
        reconstruction: [new Float32Array([1, 2, 3, 4])],
        error: 0.1,
        iterations: 10,
        converged: true,
        components: [
          {
            index: 0,
            basis: new Float32Array([1, 2, 3]),
            coefficients: new Float32Array([1, 2, 3, 4]),
            energy: 10,
            spectralCentroid: 1.5,
            spectralRolloff: 2,
            fundamental: 440,
            confidence: 0.8,
            isPitch: true,
          },
          {
            index: 1,
            basis: new Float32Array([0.1, 0.1, 0.1]),
            coefficients: new Float32Array([0.1, 0.1, 0.1, 0.1]),
            energy: 0.1,
            spectralCentroid: 1.5,
            spectralRolloff: 2,
            fundamental: 50,
            confidence: 0.2,
            isPitch: false,
          },
        ],
      };

      const pitchComponents = NMFUtils.extractPitchComponents(result);
      
      expect(pitchComponents).toHaveLength(1);
      expect(pitchComponents[0].isPitch).toBe(true);
    });
  });

  describe('groupComponentsByPitch', () => {
    it('should group components by pitch', () => {
      const components: NMFComponent[] = [
        {
          index: 0,
          basis: new Float32Array([1, 2, 3]),
          coefficients: new Float32Array([1, 2, 3, 4]),
          energy: 10,
          spectralCentroid: 1.5,
          spectralRolloff: 2,
          fundamental: 440,
          confidence: 0.8,
          isPitch: true,
        },
        {
          index: 1,
          basis: new Float32Array([2, 4, 6]),
          coefficients: new Float32Array([2, 4, 6, 8]),
          energy: 20,
          spectralCentroid: 1.5,
          spectralRolloff: 2,
          fundamental: 880,
          confidence: 0.9,
          isPitch: true,
        },
        {
          index: 2,
          basis: new Float32Array([1, 2, 3]),
          coefficients: new Float32Array([1, 2, 3, 4]),
          energy: 10,
          spectralCentroid: 1.5,
          spectralRolloff: 2,
          fundamental: 440,
          confidence: 0.7,
          isPitch: true,
        },
      ];

      const groups = NMFUtils.groupComponentsByPitch(components);
      
      expect(groups.size).toBe(2); // Two different pitches
      expect(groups.get(69)).toHaveLength(2); // Two components for A4
      expect(groups.get(81)).toHaveLength(1); // One component for A5
    });
  });

  describe('calculateSimilarity', () => {
    it('should calculate component similarity', () => {
      const comp1: NMFComponent = {
        index: 0,
        basis: new Float32Array([1, 0, 0]),
        coefficients: new Float32Array([1, 2, 3, 4]),
        energy: 10,
        spectralCentroid: 1.5,
        spectralRolloff: 2,
        fundamental: 440,
        confidence: 0.8,
        isPitch: true,
      };

      const comp2: NMFComponent = {
        index: 1,
        basis: new Float32Array([1, 0, 0]),
        coefficients: new Float32Array([2, 4, 6, 8]),
        energy: 20,
        spectralCentroid: 1.5,
        spectralRolloff: 2,
        fundamental: 880,
        confidence: 0.9,
        isPitch: true,
      };

      const similarity = NMFUtils.calculateSimilarity(comp1, comp2);
      
      expect(similarity).toBeCloseTo(1, 2); // Identical basis vectors
    });

    it('should handle orthogonal components', () => {
      const comp1: NMFComponent = {
        index: 0,
        basis: new Float32Array([1, 0, 0]),
        coefficients: new Float32Array([1, 2, 3, 4]),
        energy: 10,
        spectralCentroid: 1.5,
        spectralRolloff: 2,
        fundamental: 440,
        confidence: 0.8,
        isPitch: true,
      };

      const comp2: NMFComponent = {
        index: 1,
        basis: new Float32Array([0, 1, 0]),
        coefficients: new Float32Array([2, 4, 6, 8]),
        energy: 20,
        spectralCentroid: 1.5,
        spectralRolloff: 2,
        fundamental: 880,
        confidence: 0.9,
        isPitch: true,
      };

      const similarity = NMFUtils.calculateSimilarity(comp1, comp2);
      
      expect(similarity).toBeCloseTo(0, 2); // Orthogonal basis vectors
    });
  });
});

describe('decomposeNMF utility', () => {
  it('should decompose matrix with default config', () => {
    const spectrogram = [
      new Float32Array([1, 2, 3, 4]),
      new Float32Array([2, 4, 6, 8]),
    ];

    const result = decomposeNMF(spectrogram);
    
    expect(result.W).toHaveLength(4); // Default rank
    expect(result.H).toHaveLength(4);
    expect(result.error).toBeGreaterThanOrEqual(0);
    expect(result.components).toHaveLength(4);
  });

  it('should use custom config', () => {
    const spectrogram = [
      new Float32Array([1, 2, 3, 4]),
      new Float32Array([2, 4, 6, 8]),
    ];

    const customConfig = {
      rank: 2,
      maxIterations: 20,
      tolerance: 1e-4,
    };

    const result = decomposeNMF(spectrogram, customConfig);
    
    expect(result.W).toHaveLength(2);
    expect(result.H).toHaveLength(2);
    expect(result.components).toHaveLength(2);
  });
});