/**
 * CREPE-NMF Integration for Polyphonic Pitch Detection
 * 
 * This module combines CREPE neural network validation with NMF source separation
 * to provide accurate polyphonic pitch detection with validation.
 */

import { CREPEModel, CREPEConfig, CREPEPrediction } from './crepe';
import { NMFAlgorithm, NMFResult, NMFComponent, NMFUtils } from './nmf';
import { WindowFunction } from './dsp';

/**
 * Configuration for CREPE-NMF integration
 */
export interface CREPENMFConfig {
  crepeConfig?: CREPEConfig;
  nmfConfig?: {
    rank?: number;
    maxIterations?: number;
    tolerance?: number;
    sparsity?: number;
    smoothness?: number;
  };
  validationThreshold?: number;    // Minimum confidence for validation
  maxPitches?: number;             // Maximum number of pitches to detect
  useValidation?: boolean;         // Enable CREPE validation
  useNMF?: boolean;               // Enable NMF source separation
  sampleRate?: number;            // Input sample rate
  frameSize?: number;             // Frame size for processing
}

/**
 * Result of CREPE-NMF polyphonic detection
 */
export interface CREPENMFResult {
  pitches: CREPENMFPitch[];
  nmfResult?: NMFResult;
  crepeResults?: CREPEPrediction[];
  validationScores: number[];
  processingTime: number;
  method: 'crepe-only' | 'nmf-only' | 'crepe-nmf-hybrid';
}

/**
 * Individual pitch detection with validation
 */
export interface CREPENMFPitch {
  frequency: number;
  midi: number;
  note: string;
  confidence: number;
  clarity: number;
  timestamp: number;
  source: 'crepe' | 'nmf' | 'validated';
  validationScore: number;
  nmfComponent?: NMFComponent;
  crepePrediction?: CREPEPrediction;
}

/**
 * CREPE-NMF hybrid detector
 */
export class CREPENMFDetector {
  private crepeModel: CREPEModel;
  private nmfAlgorithm: NMFAlgorithm;
  private config: Required<CREPENMFConfig>;

  constructor(config: CREPENMFConfig = {}) {
    this.config = {
      crepeConfig: config.crepeConfig ?? {},
      nmfConfig: {
        rank: config.nmfConfig?.rank ?? 4,
        maxIterations: config.nmfConfig?.maxIterations ?? 100,
        tolerance: config.nmfConfig?.tolerance ?? 1e-6,
        sparsity: config.nmfConfig?.sparsity ?? 0.1,
        smoothness: config.nmfConfig?.smoothness ?? 0.1,
      },
      validationThreshold: config.validationThreshold ?? 0.7,
      maxPitches: config.maxPitches ?? 4,
      useValidation: config.useValidation ?? true,
      useNMF: config.useNMF ?? true,
      sampleRate: config.sampleRate ?? 44100,
      frameSize: config.frameSize ?? 2048,
    };

    this.crepeModel = new CREPEModel(this.config.crepeConfig);
    this.nmfAlgorithm = new NMFAlgorithm({
      rank: this.config.nmfConfig.rank,
      maxIterations: this.config.nmfConfig.maxIterations,
      tolerance: this.config.nmfConfig.tolerance,
      sparsity: this.config.nmfConfig.sparsity,
      smoothness: this.config.nmfConfig.smoothness,
      useMultiplicative: true,
      useAlternating: false,
      useKullbackLeibler: false,
      useEuclidean: true,
      randomSeed: 42,
    });
  }

  /**
   * Initialize the detector
   */
  async initialize(): Promise<void> {
    await this.crepeModel.initialize();
  }

  /**
   * Detect polyphonic pitches using CREPE-NMF hybrid approach
   */
  async detectPolyphonicPitches(samples: Float32Array): Promise<CREPENMFResult> {
    const startTime = performance.now();
    const result: CREPENMFResult = {
      pitches: [],
      validationScores: [],
      processingTime: 0,
      method: 'crepe-only',
    };

    try {
      // Step 1: NMF Source Separation
      let nmfResult: NMFResult | undefined;
      let nmfPitches: CREPENMFPitch[] = [];

      if (this.config.useNMF) {
        nmfResult = await this.performNMFDetection(samples);
        nmfPitches = this.extractNMFPitches(nmfResult);
        result.nmfResult = nmfResult;
      }

      // Step 2: CREPE Validation
      let crepeResults: CREPEPrediction[] = [];
      let crepePitches: CREPENMFPitch[] = [];

      if (this.config.useValidation) {
        crepeResults = await this.performCREPEDetection(samples);
        crepePitches = this.extractCREPEPitches(crepeResults);
        result.crepeResults = crepeResults;
      }

      // Step 3: Hybrid Integration and Validation
      if (this.config.useNMF && this.config.useValidation) {
        result.pitches = this.combineAndValidatePitches(nmfPitches, crepePitches);
        result.method = 'crepe-nmf-hybrid';
      } else if (this.config.useNMF) {
        result.pitches = nmfPitches;
        result.method = 'nmf-only';
      } else if (this.config.useValidation) {
        result.pitches = crepePitches;
        result.method = 'crepe-only';
      }

      // Step 4: Calculate validation scores
      result.validationScores = this.calculateValidationScores(result.pitches);

    } catch (error) {
      console.warn('CREPE-NMF detection failed:', error);
      // Fallback to basic detection
      result.pitches = await this.fallbackDetection(samples);
    }

    result.processingTime = performance.now() - startTime;
    return result;
  }

  /**
   * Perform NMF-based detection
   */
  private async performNMFDetection(samples: Float32Array): Promise<NMFResult> {
    // Create spectrogram
    const spectrogram = this.createSpectrogram(samples);
    
    // Perform NMF decomposition
    const nmfResult = this.nmfAlgorithm.decompose(spectrogram);
    
    return nmfResult;
  }

  /**
   * Extract pitches from NMF result
   */
  private extractNMFPitches(nmfResult: NMFResult): CREPENMFPitch[] {
    const pitchComponents = NMFUtils.extractPitchComponents(nmfResult);
    const pitches: CREPENMFPitch[] = [];

    for (const component of pitchComponents) {
      if (component.isPitch && component.confidence > this.config.validationThreshold) {
        const midi = Math.round(12 * Math.log2(component.fundamental / 440) + 69);
        pitches.push({
          frequency: component.fundamental,
          midi,
          note: this.midiToNoteName(midi),
          confidence: component.confidence,
          clarity: component.confidence,
          timestamp: Date.now(),
          source: 'nmf',
          validationScore: component.confidence,
          nmfComponent: component,
        });
      }
    }

    return pitches.slice(0, this.config.maxPitches);
  }

  /**
   * Perform CREPE-based detection
   */
  private async performCREPEDetection(samples: Float32Array): Promise<CREPEPrediction[]> {
    // Resample to 16kHz for CREPE
    const resampledSamples = this.resampleTo16kHz(samples);
    
    // Process with CREPE
    const prediction = await this.crepeModel.predict(resampledSamples);
    
    return [prediction];
  }

  /**
   * Extract pitches from CREPE results
   */
  private extractCREPEPitches(crepeResults: CREPEPrediction[]): CREPENMFPitch[] {
    const pitches: CREPENMFPitch[] = [];

    for (const prediction of crepeResults) {
      if (prediction.confidence > this.config.validationThreshold) {
        const midi = Math.round(12 * Math.log2(prediction.frequency / 440) + 69);
        pitches.push({
          frequency: prediction.frequency,
          midi,
          note: this.midiToNoteName(midi),
          confidence: prediction.confidence,
          clarity: prediction.confidence,
          timestamp: Date.now(),
          source: 'crepe',
          validationScore: prediction.confidence,
          crepePrediction: prediction,
        });
      }
    }

    return pitches.slice(0, this.config.maxPitches);
  }

  /**
   * Combine and validate pitches from both methods
   */
  private combineAndValidatePitches(
    nmfPitches: CREPENMFPitch[], 
    crepePitches: CREPENMFPitch[]
  ): CREPENMFPitch[] {
    const combinedPitches: CREPENMFPitch[] = [];
    const usedCrepePitches = new Set<number>();

    // Match NMF pitches with CREPE pitches
    for (const nmfPitch of nmfPitches) {
      const bestMatch = this.findBestCREPEMatch(nmfPitch, crepePitches, usedCrepePitches);
      
      if (bestMatch) {
        // Validated pitch
        combinedPitches.push({
          ...nmfPitch,
          source: 'validated',
          validationScore: (nmfPitch.validationScore + bestMatch.validationScore) / 2,
          crepePrediction: bestMatch.crepePrediction,
        });
        usedCrepePitches.add(bestMatch.midi);
      } else {
        // NMF-only pitch (lower confidence)
        combinedPitches.push({
          ...nmfPitch,
          validationScore: nmfPitch.validationScore * 0.7, // Reduce confidence
        });
      }
    }

    // Add unmatched CREPE pitches
    for (const crepePitch of crepePitches) {
      if (!usedCrepePitches.has(crepePitch.midi)) {
        combinedPitches.push({
          ...crepePitch,
          validationScore: crepePitch.validationScore * 0.8, // Slightly reduce confidence
        });
      }
    }

    // Sort by validation score and limit
    return combinedPitches
      .sort((a, b) => b.validationScore - a.validationScore)
      .slice(0, this.config.maxPitches);
  }

  /**
   * Find best CREPE match for NMF pitch
   */
  private findBestCREPEMatch(
    nmfPitch: CREPENMFPitch, 
    crepePitches: CREPENMFPitch[], 
    usedCrepePitches: Set<number>
  ): CREPENMFPitch | null {
    let bestMatch: CREPENMFPitch | null = null;
    let bestScore = 0;

    for (const crepePitch of crepePitches) {
      if (usedCrepePitches.has(crepePitch.midi)) continue;

      // Calculate matching score
      const frequencyDiff = Math.abs(nmfPitch.frequency - crepePitch.frequency);
      const frequencyScore = Math.exp(-frequencyDiff / 50); // 50 Hz tolerance
      
      const confidenceScore = Math.min(nmfPitch.confidence, crepePitch.confidence);
      const totalScore = frequencyScore * confidenceScore;

      if (totalScore > bestScore && totalScore > 0.5) {
        bestMatch = crepePitch;
        bestScore = totalScore;
      }
    }

    return bestMatch;
  }

  /**
   * Calculate validation scores for all pitches
   */
  private calculateValidationScores(pitches: CREPENMFPitch[]): number[] {
    return pitches.map(pitch => pitch.validationScore);
  }

  /**
   * Fallback detection when hybrid approach fails
   */
  private async fallbackDetection(samples: Float32Array): Promise<CREPENMFPitch[]> {
    // Simple fallback - just use CREPE if available
    if (this.config.useValidation) {
      try {
        const crepeResults = await this.performCREPEDetection(samples);
        return this.extractCREPEPitches(crepeResults);
      } catch (error) {
        console.warn('Fallback CREPE detection failed:', error);
      }
    }

    return [];
  }

  /**
   * Create spectrogram from samples
   */
  private createSpectrogram(samples: Float32Array): Float32Array[] {
    const hopSize = this.config.frameSize / 2;
    const frameSize = this.config.frameSize;
    const spectrogram: Float32Array[] = [];

    // For single frame processing
    if (samples.length === frameSize) {
      const windowedFrame = this.applyWindow(samples);
      const magnitude = this.computeFFTMagnitude(windowedFrame);
      spectrogram.push(magnitude);
    } else {
      // Extract overlapping frames
      for (let i = 0; i <= samples.length - frameSize; i += hopSize) {
        const frame = samples.slice(i, i + frameSize);
        const windowedFrame = this.applyWindow(frame);
        const magnitude = this.computeFFTMagnitude(windowedFrame);
        spectrogram.push(magnitude);
      }
    }

    return spectrogram;
  }

  /**
   * Apply window function
   */
  private applyWindow(frame: Float32Array): Float32Array {
    const window = WindowFunction.create(frame.length, 'hann');
    return WindowFunction.apply(frame, window);
  }

  /**
   * Compute FFT magnitude (simplified)
   */
  private computeFFTMagnitude(frame: Float32Array): Float32Array {
    const magnitude = new Float32Array(frame.length / 2 + 1);
    for (let i = 0; i < magnitude.length; i++) {
      magnitude[i] = Math.abs(frame[i] || 0);
    }
    return magnitude;
  }

  /**
   * Resample to 16kHz for CREPE
   */
  private resampleTo16kHz(samples: Float32Array): Float32Array {
    const targetSampleRate = 16000;
    const ratio = this.config.sampleRate / targetSampleRate;
    const targetLength = Math.floor(samples.length / ratio);
    const resampled = new Float32Array(targetLength);

    for (let i = 0; i < targetLength; i++) {
      const sourceIndex = i * ratio;
      const index = Math.floor(sourceIndex);
      const fraction = sourceIndex - index;
      
      if (index < samples.length - 1) {
        resampled[i] = samples[index] * (1 - fraction) + samples[index + 1] * fraction;
      } else {
        resampled[i] = samples[index] || 0;
      }
    }

    return resampled;
  }

  /**
   * Convert MIDI to note name
   */
  private midiToNoteName(midi: number): string {
    const notes = ['C', 'C#', 'D', 'D#', 'E', 'F', 'F#', 'G', 'G#', 'A', 'A#', 'B'];
    const octave = Math.floor(midi / 12) - 1;
    const note = notes[midi % 12];
    return `${note}${octave}`;
  }

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

  /**
   * Update configuration
   */
  updateConfig(newConfig: Partial<CREPENMFConfig>): void {
    this.config = { ...this.config, ...newConfig };
    
    if (newConfig.nmfConfig) {
      this.nmfAlgorithm.updateConfig(newConfig.nmfConfig);
    }
  }

  /**
   * Dispose resources
   */
  dispose(): void {
    this.crepeModel.dispose();
  }
}

/**
 * Quick utility function for CREPE-NMF detection
 */
export async function detectPolyphonicPitches(
  samples: Float32Array,
  config: CREPENMFConfig = {}
): Promise<CREPENMFResult> {
  const detector = new CREPENMFDetector(config);
  await detector.initialize();
  const result = await detector.detectPolyphonicPitches(samples);
  detector.dispose();
  return result;
}