/**
 * Non-negative Matrix Factorization (NMF) for polyphonic pitch detection
 * 
 * This module implements NMF algorithms for source separation and spectral
 * dictionary learning to enable detection of multiple simultaneous pitches.
 */

/**
 * Configuration for NMF algorithm
 */
export interface NMFConfig {
  rank?: number;              // Number of components (default: 4)
  maxIterations?: number;     // Maximum iterations (default: 100)
  tolerance?: number;         // Convergence tolerance (default: 1e-6)
  sparsity?: number;          // Sparsity regularization (default: 0.1)
  smoothness?: number;        // Smoothness regularization (default: 0.1)
  useMultiplicative?: boolean; // Use multiplicative updates (default: true)
  useAlternating?: boolean;   // Use alternating least squares (default: false)
  useKullbackLeibler?: boolean; // Use KL divergence (default: false)
  useEuclidean?: boolean;     // Use Euclidean distance (default: true)
  randomSeed?: number;        // Random seed for initialization (default: 42)
}

/**
 * NMF result containing factorized matrices
 */
export interface NMFResult {
  W: Float32Array[];          // Basis matrix (rank x features)
  H: Float32Array[];          // Coefficient matrix (rank x time)
  reconstruction: Float32Array[]; // Reconstructed matrix
  error: number;              // Final reconstruction error
  iterations: number;          // Number of iterations performed
  converged: boolean;         // Whether algorithm converged
  components: NMFComponent[]; // Individual components with metadata
}

/**
 * Individual NMF component with metadata
 */
export interface NMFComponent {
  index: number;              // Component index
  basis: Float32Array;        // Basis vector (spectral pattern)
  coefficients: Float32Array; // Coefficient vector (time series)
  energy: number;            // Total energy of component
  spectralCentroid: number;   // Spectral centroid of basis
  spectralRolloff: number;    // Spectral rolloff of basis
  fundamental: number;        // Estimated fundamental frequency
  confidence: number;         // Component confidence (0-1)
  isPitch: boolean;          // Whether component represents a pitch
}

/**
 * Spectral dictionary for NMF
 */
export interface SpectralDictionary {
  atoms: Float32Array[];      // Dictionary atoms (spectral patterns)
  frequencies: number[];      // Corresponding frequencies
  pitches: number[];          // Corresponding MIDI pitches
  weights: number[];          // Atom weights
  size: number;               // Dictionary size
}

/**
 * Main NMF algorithm class
 */
export class NMFAlgorithm {
  private config: Required<NMFConfig>;
  private random: () => number;

  constructor(config: NMFConfig = {}) {
    this.config = {
      rank: config.rank ?? 4,
      maxIterations: config.maxIterations ?? 100,
      tolerance: config.tolerance ?? 1e-6,
      sparsity: config.sparsity ?? 0.1,
      smoothness: config.smoothness ?? 0.1,
      useMultiplicative: config.useMultiplicative ?? true,
      useAlternating: config.useAlternating ?? false,
      useKullbackLeibler: config.useKullbackLeibler ?? false,
      useEuclidean: config.useEuclidean ?? true,
      randomSeed: config.randomSeed ?? 42,
    };

    // Initialize random number generator
    this.random = this.createRandomGenerator(this.config.randomSeed);
  }

  /**
   * Perform NMF decomposition
   */
  decompose(
    V: Float32Array[], 
    initialW?: Float32Array[], 
    initialH?: Float32Array[]
  ): NMFResult {
    if (V.length === 0 || V[0].length === 0) {
      throw new Error('Input matrix V cannot be empty');
    }
    
    const [rows, cols] = [V.length, V[0].length];
    const rank = this.config.rank;

    // Initialize matrices
    let W = initialW || this.initializeW(rows, rank);
    let H = initialH || this.initializeH(rank, cols);

    let error = Infinity;
    let prevError = Infinity;
    let iterations = 0;
    let converged = false;

    // Main iteration loop
    for (let iter = 0; iter < this.config.maxIterations; iter++) {
      iterations = iter + 1;

      if (this.config.useMultiplicative) {
        this.multiplicativeUpdate(V, W, H);
      } else if (this.config.useAlternating) {
        this.alternatingLeastSquares(V, W, H);
      }

      // Calculate reconstruction error
      prevError = error;
      error = this.calculateError(V, W, H);

      // Check convergence
      if (Math.abs(prevError - error) < this.config.tolerance) {
        converged = true;
        break;
      }
    }

    // Create reconstruction
    const reconstruction = this.reconstruct(W, H);

    // Extract components
    const components = this.extractComponents(W, H);

    return {
      W,
      H,
      reconstruction,
      error,
      iterations,
      converged,
      components,
    };
  }

  /**
   * Initialize basis matrix W
   */
  private initializeW(rows: number, rank: number): Float32Array[] {
    const W: Float32Array[] = [];
    for (let r = 0; r < rank; r++) {
      const basis = new Float32Array(rows);
      for (let i = 0; i < rows; i++) {
        basis[i] = this.random() * 0.1 + 0.01; // Small positive values
      }
      W.push(basis);
    }
    return W;
  }

  /**
   * Initialize coefficient matrix H
   */
  private initializeH(rank: number, cols: number): Float32Array[] {
    const H: Float32Array[] = [];
    for (let r = 0; r < rank; r++) {
      const coeffs = new Float32Array(cols);
      for (let j = 0; j < cols; j++) {
        coeffs[j] = this.random() * 0.1 + 0.01; // Small positive values
      }
      H.push(coeffs);
    }
    return H;
  }

  /**
   * Multiplicative update rules
   */
  private multiplicativeUpdate(V: Float32Array[], W: Float32Array[], H: Float32Array[]): void {
    const [rows, cols] = [V.length, V[0].length];
    const rank = W.length;

    // Update H (coefficients)
    for (let r = 0; r < rank; r++) {
      for (let j = 0; j < cols; j++) {
        let numerator = 0;
        let denominator = 0;

        for (let i = 0; i < rows; i++) {
          const WH = this.calculateWH(W, H, i, j);
          numerator += W[r][i] * V[i][j] / (WH + 1e-10);
          denominator += W[r][i];
        }

        H[r][j] = H[r][j] * numerator / (denominator + 1e-10);
      }
    }

    // Update W (basis)
    for (let r = 0; r < rank; r++) {
      for (let i = 0; i < rows; i++) {
        let numerator = 0;
        let denominator = 0;

        for (let j = 0; j < cols; j++) {
          const WH = this.calculateWH(W, H, i, j);
          numerator += H[r][j] * V[i][j] / (WH + 1e-10);
          denominator += H[r][j];
        }

        W[r][i] = W[r][i] * numerator / (denominator + 1e-10);
      }
    }
  }

  /**
   * Alternating least squares update
   */
  private alternatingLeastSquares(V: Float32Array[], W: Float32Array[], H: Float32Array[]): void {
    const [rows, cols] = [V.length, V[0].length];
    const rank = W.length;

    // Update H using least squares
    for (let r = 0; r < rank; r++) {
      for (let j = 0; j < cols; j++) {
        let numerator = 0;
        let denominator = 0;

        for (let i = 0; i < rows; i++) {
          const WH = this.calculateWH(W, H, i, j);
          numerator += W[r][i] * (V[i][j] - WH + W[r][i] * H[r][j]);
          denominator += W[r][i] * W[r][i];
        }

        H[r][j] = Math.max(0, numerator / (denominator + 1e-10));
      }
    }

    // Update W using least squares
    for (let r = 0; r < rank; r++) {
      for (let i = 0; i < rows; i++) {
        let numerator = 0;
        let denominator = 0;

        for (let j = 0; j < cols; j++) {
          const WH = this.calculateWH(W, H, i, j);
          numerator += H[r][j] * (V[i][j] - WH + W[r][i] * H[r][j]);
          denominator += H[r][j] * H[r][j];
        }

        W[r][i] = Math.max(0, numerator / (denominator + 1e-10));
      }
    }
  }

  /**
   * Calculate WH product for element (i, j)
   */
  private calculateWH(W: Float32Array[], H: Float32Array[], i: number, j: number): number {
    let sum = 0;
    for (let r = 0; r < W.length; r++) {
      sum += W[r][i] * H[r][j];
    }
    return sum;
  }

  /**
   * Calculate reconstruction error
   */
  private calculateError(V: Float32Array[], W: Float32Array[], H: Float32Array[]): number {
    const [rows, cols] = [V.length, V[0].length];
    let error = 0;

    for (let i = 0; i < rows; i++) {
      for (let j = 0; j < cols; j++) {
        const WH = this.calculateWH(W, H, i, j);
        const diff = V[i][j] - WH;
        
        if (this.config.useKullbackLeibler) {
          // KL divergence
          error += V[i][j] * Math.log((V[i][j] + 1e-10) / (WH + 1e-10)) - V[i][j] + WH;
        } else {
          // Euclidean distance
          error += diff * diff;
        }
      }
    }

    return error;
  }

  /**
   * Reconstruct matrix from factors
   */
  private reconstruct(W: Float32Array[], H: Float32Array[]): Float32Array[] {
    const [rows, cols] = [W[0].length, H[0].length];
    const reconstruction: Float32Array[] = [];

    for (let i = 0; i < rows; i++) {
      const row = new Float32Array(cols);
      for (let j = 0; j < cols; j++) {
        row[j] = this.calculateWH(W, H, i, j);
      }
      reconstruction.push(row);
    }

    return reconstruction;
  }

  /**
   * Extract components with metadata
   */
  private extractComponents(W: Float32Array[], H: Float32Array[]): NMFComponent[] {
    const components: NMFComponent[] = [];

    for (let r = 0; r < W.length; r++) {
      const basis = W[r];
      const coefficients = H[r];
      
      // Calculate component energy
      const energy = coefficients.reduce((sum, val) => sum + val * val, 0);
      
      // Calculate spectral features
      const spectralCentroid = this.calculateSpectralCentroid(basis);
      const spectralRolloff = this.calculateSpectralRolloff(basis);
      
      // Estimate fundamental frequency
      const fundamental = this.estimateFundamental(basis);
      
      // Calculate confidence
      const confidence = this.calculateConfidence(basis, coefficients);
      
      // Determine if component represents a pitch
      const isPitch = this.isPitchComponent(basis, fundamental, confidence);

      components.push({
        index: r,
        basis,
        coefficients,
        energy,
        spectralCentroid,
        spectralRolloff,
        fundamental,
        confidence,
        isPitch,
      });
    }

    return components;
  }

  /**
   * Calculate spectral centroid
   */
  private calculateSpectralCentroid(spectrum: Float32Array): number {
    let weightedSum = 0;
    let magnitudeSum = 0;

    for (let i = 0; i < spectrum.length; i++) {
      weightedSum += i * spectrum[i];
      magnitudeSum += spectrum[i];
    }

    return magnitudeSum > 0 ? weightedSum / magnitudeSum : 0;
  }

  /**
   * Calculate spectral rolloff
   */
  private calculateSpectralRolloff(spectrum: Float32Array): number {
    const totalEnergy = spectrum.reduce((sum, val) => sum + val * val, 0);
    const threshold = 0.85 * totalEnergy;
    
    let cumulativeEnergy = 0;
    for (let i = 0; i < spectrum.length; i++) {
      cumulativeEnergy += spectrum[i] * spectrum[i];
      if (cumulativeEnergy >= threshold) {
        return i;
      }
    }

    return spectrum.length - 1;
  }

  /**
   * Estimate fundamental frequency from basis
   */
  private estimateFundamental(basis: Float32Array): number {
    // Find peak in spectrum
    let maxIndex = 0;
    let maxValue = 0;

    for (let i = 0; i < basis.length; i++) {
      if (basis[i] > maxValue) {
        maxValue = basis[i];
        maxIndex = i;
      }
    }

    // Convert bin index to frequency (assuming 44100 Hz sample rate)
    return (maxIndex * 44100) / (2 * basis.length);
  }

  /**
   * Calculate component confidence
   */
  private calculateConfidence(basis: Float32Array, coefficients: Float32Array): number {
    // Calculate sparsity of basis
    const basisSparsity = this.calculateSparsity(basis);
    
    // Calculate consistency of coefficients
    const coefficientConsistency = this.calculateConsistency(coefficients);
    
    // Combine metrics
    return (basisSparsity + coefficientConsistency) / 2;
  }

  /**
   * Calculate sparsity of a vector
   */
  private calculateSparsity(vector: Float32Array): number {
    const l1 = vector.reduce((sum, val) => sum + Math.abs(val), 0);
    const l2 = Math.sqrt(vector.reduce((sum, val) => sum + val * val, 0));
    return l2 > 0 ? 1 - (l1 / l2) : 0;
  }

  /**
   * Calculate consistency of coefficients
   */
  private calculateConsistency(coefficients: Float32Array): number {
    const mean = coefficients.reduce((sum, val) => sum + val, 0) / coefficients.length;
    const variance = coefficients.reduce((sum, val) => sum + (val - mean) ** 2, 0) / coefficients.length;
    const stdDev = Math.sqrt(variance);
    
    return stdDev > 0 ? 1 / (1 + stdDev / mean) : 1;
  }

  /**
   * Determine if component represents a pitch
   */
  private isPitchComponent(basis: Float32Array, fundamental: number, confidence: number): boolean {
    // Check if fundamental is in valid range
    const validFrequency = fundamental >= 50 && fundamental <= 2000;
    
    // Check if confidence is high enough
    const highConfidence = confidence >= 0.5;
    
    // Check if basis has clear spectral structure
    const clearStructure = this.hasClearSpectralStructure(basis);
    
    return validFrequency && highConfidence && clearStructure;
  }

  /**
   * Check if basis has clear spectral structure
   */
  private hasClearSpectralStructure(basis: Float32Array): boolean {
    // Find peaks in spectrum
    const peaks = this.findPeaks(basis);
    
    // Check if there's a dominant peak
    if (peaks.length === 0) return false;
    
    const maxPeak = Math.max(...peaks.map(p => p.strength));
    const totalEnergy = basis.reduce((sum, val) => sum + val * val, 0);
    
    return maxPeak > 0.3 * totalEnergy;
  }

  /**
   * Find peaks in spectrum
   */
  private findPeaks(spectrum: Float32Array): Array<{ index: number; strength: number }> {
    const peaks: Array<{ index: number; strength: number }> = [];
    const threshold = 0.1 * Math.max(...spectrum);
    
    for (let i = 1; i < spectrum.length - 1; i++) {
      if (spectrum[i] > threshold && 
          spectrum[i] > spectrum[i - 1] && 
          spectrum[i] > spectrum[i + 1]) {
        peaks.push({ index: i, strength: spectrum[i] });
      }
    }
    
    return peaks;
  }

  /**
   * Create random number generator
   */
  private createRandomGenerator(seed: number): () => number {
    let state = seed;
    return () => {
      state = (state * 1664525 + 1013904223) % 4294967296;
      return state / 4294967296;
    };
  }

  /**
   * Get current configuration
   */
  getConfig(): Readonly<Required<NMFConfig>> {
    return { ...this.config };
  }

  /**
   * Update configuration
   */
  updateConfig(newConfig: Partial<NMFConfig>): void {
    this.config = { ...this.config, ...newConfig };
    
    // Recreate random generator if seed changed
    if (newConfig.randomSeed !== undefined) {
      this.random = this.createRandomGenerator(this.config.randomSeed);
    }
  }
}

/**
 * Spectral dictionary learning utilities
 */
export const SpectralDictionary = {
  /**
   * Learn dictionary from training data
   */
  learnDictionary(
    trainingData: Float32Array[][],
    dictionarySize: number = 50,
    config: NMFConfig = {}
  ): SpectralDictionary {
    // Flatten training data
    const flattenedData = this.flattenTrainingData(trainingData);
    
    if (flattenedData.length === 0) {
      // Return empty dictionary if no training data
      return {
        atoms: [],
        frequencies: [],
        pitches: [],
        weights: [],
        size: 0,
      };
    }
    
    // Perform NMF with higher rank
    const nmf = new NMFAlgorithm({ ...config, rank: dictionarySize });
    const result = nmf.decompose(flattenedData);
    
    // Extract dictionary atoms
    const atoms = result.W;
    const frequencies = atoms.map(atom => this.estimateFrequency(atom));
    const pitches = frequencies.map(freq => {
      const pitch = Math.round(12 * Math.log2(freq / 440) + 69);
      return Math.max(0, Math.min(127, pitch)); // Clamp to valid MIDI range
    });
    const weights = atoms.map(atom => this.calculateAtomWeight(atom));
    
    return {
      atoms,
      frequencies,
      pitches,
      weights,
      size: atoms.length,
    };
  },

  /**
   * Flatten training data
   */
  flattenTrainingData(trainingData: Float32Array[][]): Float32Array[] {
    const flattened: Float32Array[] = [];
    
    for (const sample of trainingData) {
      for (const frame of sample) {
        flattened.push(frame);
      }
    }
    
    return flattened;
  },

  /**
   * Estimate frequency from atom
   */
  estimateFrequency(atom: Float32Array): number {
    let maxIndex = 0;
    let maxValue = 0;
    
    for (let i = 0; i < atom.length; i++) {
      if (atom[i] > maxValue) {
        maxValue = atom[i];
        maxIndex = i;
      }
    }
    
    return (maxIndex * 44100) / (2 * atom.length);
  },

  /**
   * Calculate atom weight
   */
  calculateAtomWeight(atom: Float32Array): number {
    return atom.reduce((sum, val) => sum + val * val, 0);
  },
};

/**
 * NMF utilities for polyphonic detection
 */
export const NMFUtils = {
  /**
   * Separate sources using NMF
   */
  separateSources(
    spectrogram: Float32Array[],
    dictionary: SpectralDictionary,
    config: NMFConfig = {}
  ): NMFResult {
    const nmf = new NMFAlgorithm({ ...config, rank: dictionary.size });
    
    // Initialize with dictionary
    const initialW = dictionary.atoms;
    const result = nmf.decompose(spectrogram, initialW);
    
    return result;
  },

  /**
   * Extract pitch components
   */
  extractPitchComponents(result: NMFResult): NMFComponent[] {
    return result.components.filter(comp => comp.isPitch);
  },

  /**
   * Group components by pitch
   */
  groupComponentsByPitch(components: NMFComponent[]): Map<number, NMFComponent[]> {
    const groups = new Map<number, NMFComponent[]>();
    
    for (const component of components) {
      const pitch = Math.round(12 * Math.log2(component.fundamental / 440) + 69);
      
      if (!groups.has(pitch)) {
        groups.set(pitch, []);
      }
      groups.get(pitch)!.push(component);
    }
    
    return groups;
  },

  /**
   * Calculate component similarity
   */
  calculateSimilarity(comp1: NMFComponent, comp2: NMFComponent): number {
    // Calculate cosine similarity between basis vectors
    const dotProduct = comp1.basis.reduce((sum, val, i) => sum + val * comp2.basis[i], 0);
    const norm1 = Math.sqrt(comp1.basis.reduce((sum, val) => sum + val * val, 0));
    const norm2 = Math.sqrt(comp2.basis.reduce((sum, val) => sum + val * val, 0));
    
    return norm1 > 0 && norm2 > 0 ? dotProduct / (norm1 * norm2) : 0;
  },
};

/**
 * Quick utility function for NMF decomposition
 */
export function decomposeNMF(
  spectrogram: Float32Array[],
  config: NMFConfig = {}
): NMFResult {
  const nmf = new NMFAlgorithm(config);
  return nmf.decompose(spectrogram);
}