/**
 * Polyphonic Pitch Detection Package
 * Uses CREPE-inspired neural networks with NMF source separation
 */

import { AudioProcessor } from './audio';
import { AutocorrelationDetector, detectPitchAutocorrelation, AutocorrelationUtils } from './autocorrelation';
import { HarmonicAnalyzer, analyzeHarmonics } from './harmonic-analysis';
import { PitchTracker } from './pitch-tracking';
import { NMFAlgorithm, SpectralDictionary, NMFUtils } from './nmf';
import { CREPENMFDetector, detectPolyphonicPitches } from './crepe-nmf';

/**
 * Configuration options for the pitch detector
 */
export interface PitchDetectionOptions {
  sampleRate?: number;          // Default: 44100
  frameSize?: number;            // Default: 2048
  hopSize?: number;              // Default: 1024
  maxPolyphony?: number;         // Default: 4 (1-6 supported)
  confidenceThreshold?: number;  // Default: 0.7 (0-1)
  useNMF?: boolean;             // Default: true (use NMF for polyphonic detection)
  useCrepe?: boolean;           // Default: true
  useWorklet?: boolean;         // Default: false (use AudioWorklet for better performance)
  useHarmonicAnalysis?: boolean; // Default: true (use harmonic analysis for better accuracy)
  usePitchTracking?: boolean;   // Default: true (use pitch tracking for temporal smoothing)
  nmfRank?: number;             // Default: 4 (NMF rank for polyphonic detection)
  useCREPENMF?: boolean;        // Default: true (use CREPE-NMF hybrid for polyphonic detection)
  crepeNMFConfig?: {            // CREPE-NMF specific configuration
    validationThreshold?: number;
    maxPitches?: number;
    useValidation?: boolean;
    useNMF?: boolean;
  };
  modelPath?: string;           // Path to CREPE model
  workletPath?: string;         // Path to worklet processor
}

/**
 * Represents a detected pitch with metadata
 */
export interface DetectedPitch {
  frequency: number;            // Hz
  midi: number;                 // MIDI note number
  note: string;                 // e.g., "A4"
  confidence: number;           // 0-1
  clarity: number;              // Harmonic clarity 0-1
  timestamp: number;            // Relative to audio start (ms)
}

/**
 * Main pitch detection class
 */
export class PitchDetector {
  private options: Required<PitchDetectionOptions>;
  private initialized: boolean = false;
  private audioProcessor: AudioProcessor;
  private harmonicAnalyzer: HarmonicAnalyzer;
  private pitchTracker: PitchTracker;
  private nmfAlgorithm: NMFAlgorithm;
  private crepeNMFDetector: CREPENMFDetector;

  constructor(options: PitchDetectionOptions = {}) {
    this.options = {
      sampleRate: options.sampleRate ?? 44100,
      frameSize: options.frameSize ?? 2048,
      hopSize: options.hopSize ?? 1024,
      maxPolyphony: options.maxPolyphony ?? 4,
      confidenceThreshold: options.confidenceThreshold ?? 0.7,
      useNMF: options.useNMF ?? true,
      useCrepe: options.useCrepe ?? true,
      useWorklet: options.useWorklet ?? false,
      useHarmonicAnalysis: options.useHarmonicAnalysis ?? true,
      usePitchTracking: options.usePitchTracking ?? true,
      nmfRank: options.nmfRank ?? 4,
      useCREPENMF: options.useCREPENMF ?? true,
      crepeNMFConfig: {
        validationThreshold: options.crepeNMFConfig?.validationThreshold ?? 0.7,
        maxPitches: options.crepeNMFConfig?.maxPitches ?? 4,
        useValidation: options.crepeNMFConfig?.useValidation ?? true,
        useNMF: options.crepeNMFConfig?.useNMF ?? true,
      },
      modelPath: options.modelPath ?? '/models/crepe.json',
      workletPath: options.workletPath ?? '/src/pitch-worklet.js',
    };

    this.audioProcessor = new AudioProcessor({
      sampleRate: this.options.sampleRate,
      frameSize: this.options.frameSize,
      hopSize: this.options.hopSize,
      useWorklet: this.options.useWorklet,
      workletPath: this.options.workletPath,
    });

    this.harmonicAnalyzer = new HarmonicAnalyzer({
      sampleRate: this.options.sampleRate,
      fftSize: this.options.frameSize,
    });

    this.pitchTracker = new PitchTracker({
      sampleRate: this.options.sampleRate,
      frameSize: this.options.frameSize,
      hopSize: this.options.hopSize,
    });

    this.nmfAlgorithm = new NMFAlgorithm({
      rank: this.options.nmfRank,
      maxIterations: 100,
      tolerance: 1e-6,
      sparsity: 0.1,
      smoothness: 0.1,
      useMultiplicative: true,
      useAlternating: false,
      useKullbackLeibler: false,
      useEuclidean: true,
      randomSeed: 42,
    });

    this.crepeNMFDetector = new CREPENMFDetector({
      validationThreshold: this.options.crepeNMFConfig.validationThreshold,
      maxPitches: this.options.crepeNMFConfig.maxPitches,
      useValidation: this.options.crepeNMFConfig.useValidation,
      useNMF: this.options.crepeNMFConfig.useNMF,
      sampleRate: this.options.sampleRate,
      frameSize: this.options.frameSize,
    });
  }

  /**
   * Initialize the pitch detector (load models, setup processing)
   */
  async initialize(): Promise<void> {
    if (this.initialized) {
      return;
    }

    // Initialize audio processor
    await this.audioProcessor.initialize();

    // Initialize CREPE-NMF detector if enabled
    if (this.options.useCREPENMF) {
      try {
        await this.crepeNMFDetector.initialize();
      } catch (error) {
        console.warn('CREPE-NMF initialization failed, falling back to NMF-only:', error);
        // Disable CREPE-NMF if initialization fails
        this.options.useCREPENMF = false;
      }
    }

    this.initialized = true;
  }

  /**
   * Detect pitches from an AudioBuffer
   */
  async detectFromAudioBuffer(buffer: AudioBuffer): Promise<DetectedPitch[]> {
    if (!this.initialized) {
      throw new Error('PitchDetector must be initialized before use');
    }

    // Resample if needed
    const targetSampleRate = this.options.sampleRate;
    const resampledBuffer = await this.audioProcessor.resampleBuffer(buffer, targetSampleRate);

    // Extract frames
    const frames = this.audioProcessor.extractFrames(resampledBuffer);

    // Process all frames and collect results
    const allPitches: DetectedPitch[] = [];
    
    for (let i = 0; i < frames.length; i++) {
      const pitches = await this.processFrame(frames[i]);
      
      // Add timestamp to each pitch
      const timestamp = (i * this.options.hopSize / this.options.sampleRate) * 1000;
      pitches.forEach(pitch => {
        pitch.timestamp = timestamp;
      });
      
      allPitches.push(...pitches);
    }

    return allPitches;
  }

  /**
   * Detect pitches from a MediaStream (real-time)
   */
  async *detectFromStream(stream: MediaStream): AsyncGenerator<DetectedPitch[]> {
    if (!this.initialized) {
      throw new Error('PitchDetector must be initialized before use');
    }

    let frameCount = 0;

    // Process audio stream
    for await (const frame of this.audioProcessor.processStream(stream)) {
      const pitches = await this.processFrame(frame);
      
      // Add timestamp
      const timestamp = (frameCount * this.options.hopSize / this.options.sampleRate) * 1000;
      pitches.forEach(pitch => {
        pitch.timestamp = timestamp;
      });
      
      frameCount++;
      yield pitches;
    }
  }

  /**
   * Process a single audio frame
   */
  async processFrame(samples: Float32Array): Promise<DetectedPitch[]> {
    if (!this.initialized) {
      throw new Error('PitchDetector must be initialized before use');
    }

    if (samples.length !== this.options.frameSize) {
      throw new Error(`Frame size must be ${this.options.frameSize}, got ${samples.length}`);
    }

      // Try harmonic analysis first if enabled
      if (this.options.useHarmonicAnalysis) {
        try {
          const harmonicResult = this.harmonicAnalyzer.analyzeHarmonics(samples);
          if (harmonicResult && harmonicResult.confidence > (this.options.confidenceThreshold || 0.7)) {
            const detectedPitch = {
              frequency: harmonicResult.fundamental,
              confidence: harmonicResult.confidence,
              clarity: harmonicResult.harmonicity,
              timestamp: Date.now(),
            };

            // Apply pitch tracking if enabled
            if (this.options.usePitchTracking) {
              const trackedPitch = this.pitchTracker.processPitch(detectedPitch);
              if (trackedPitch) {
                const midi = Math.round(12 * Math.log2(trackedPitch.frequency / 440) + 69);
                return [{
                  frequency: trackedPitch.frequency,
                  midi,
                  note: midiToNoteName(midi),
                  confidence: trackedPitch.confidence,
                  clarity: trackedPitch.clarity,
                  timestamp: trackedPitch.timestamp,
                }];
              }
            } else {
              const midi = Math.round(12 * Math.log2(harmonicResult.fundamental / 440) + 69);
              return [{
                frequency: harmonicResult.fundamental,
                midi,
                note: midiToNoteName(midi),
                confidence: harmonicResult.confidence,
                clarity: harmonicResult.harmonicity,
                timestamp: 0,
              }];
            }
          }
        } catch (error) {
          // Fall back to autocorrelation if harmonic analysis fails
          console.warn('Harmonic analysis failed, falling back to autocorrelation:', error);
        }
      }

    // Use autocorrelation as fallback for monophonic detection
    const autocorrDetector = new AutocorrelationDetector({
      sampleRate: this.options.sampleRate,
      minFrequency: 50,
      maxFrequency: 2000,
      threshold: this.options.confidenceThreshold || 0.7,
    });
    
    const result = autocorrDetector.detectPitch(samples);
    
    if (result && result.confidence > (this.options.confidenceThreshold || 0.7)) {
      const detectedPitch = {
        frequency: result.frequency,
        confidence: result.confidence,
        clarity: result.confidence, // Use confidence as clarity for now
        timestamp: Date.now(),
      };

      // Apply pitch tracking if enabled
      if (this.options.usePitchTracking) {
        const trackedPitch = this.pitchTracker.processPitch(detectedPitch);
        if (trackedPitch) {
          const midi = Math.round(12 * Math.log2(trackedPitch.frequency / 440) + 69);
          return [{
            frequency: trackedPitch.frequency,
            midi,
            note: midiToNoteName(midi),
            confidence: trackedPitch.confidence,
            clarity: trackedPitch.clarity,
            timestamp: trackedPitch.timestamp,
          }];
        }
      } else {
        const midi = Math.round(12 * Math.log2(result.frequency / 440) + 69);
        return [{
          frequency: result.frequency,
          midi,
          note: midiToNoteName(midi),
          confidence: result.confidence,
          clarity: result.confidence, // Use confidence as clarity for now
          timestamp: 0,
        }];
      }
    }

    // Try CREPE-NMF hybrid for polyphonic detection if enabled
    if (this.options.useCREPENMF && this.options.maxPolyphony > 1) {
      try {
        const crepeNMFResult = await this.crepeNMFDetector.detectPolyphonicPitches(samples);
        if (crepeNMFResult.pitches.length > 0) {
          // Convert CREPE-NMF pitches to DetectedPitch format
          return crepeNMFResult.pitches.map(pitch => ({
            frequency: pitch.frequency,
            midi: pitch.midi,
            note: pitch.note,
            confidence: pitch.confidence,
            clarity: pitch.clarity,
            timestamp: pitch.timestamp,
          }));
        }
      } catch (error) {
        console.warn('CREPE-NMF polyphonic detection failed:', error);
      }
    }

    // Fallback to NMF-only for polyphonic detection if enabled
    if (this.options.useNMF && this.options.maxPolyphony > 1) {
      try {
        const polyphonicPitches = await this.detectPolyphonicPitches(samples);
        if (polyphonicPitches.length > 0) {
          return polyphonicPitches;
        }
      } catch (error) {
        console.warn('NMF polyphonic detection failed:', error);
      }
    }

    // TODO: Implement CREPE inference

    return [];
  }

  /**
   * Clean up resources
   */
  dispose(): void {
    this.audioProcessor.dispose();
    
    // Dispose CREPE-NMF detector if enabled
    if (this.options.useCREPENMF) {
      this.crepeNMFDetector.dispose();
    }
    
    this.initialized = false;
  }

  /**
   * Get the audio processor instance
   */
  getAudioProcessor(): AudioProcessor {
    return this.audioProcessor;
  }

  /**
   * Get current configuration
   */
  getOptions(): Readonly<Required<PitchDetectionOptions>> {
    return { ...this.options };
  }

  /**
   * Get pitch tracking state
   */
  getPitchTrackingState() {
    return this.pitchTracker.getState();
  }

  /**
   * Reset pitch tracking state
   */
  resetPitchTracking() {
    this.pitchTracker.reset();
  }

  /**
   * Detect polyphonic pitches using NMF
   */
  private async detectPolyphonicPitches(samples: Float32Array): Promise<DetectedPitch[]> {
    // Convert time-domain signal to spectrogram
    const spectrogram = this.createSpectrogram(samples);
    
    // Perform NMF decomposition
    const nmfResult = this.nmfAlgorithm.decompose(spectrogram);
    
    // Extract pitch components
    const pitchComponents = NMFUtils.extractPitchComponents(nmfResult);
    
    // Convert components to detected pitches
    const detectedPitches: DetectedPitch[] = [];
    
    for (const component of pitchComponents) {
      if (component.isPitch && component.confidence > (this.options.confidenceThreshold || 0.7)) {
        const midi = Math.round(12 * Math.log2(component.fundamental / 440) + 69);
        detectedPitches.push({
          frequency: component.fundamental,
          midi,
          note: midiToNoteName(midi),
          confidence: component.confidence,
          clarity: component.confidence, // Use confidence as clarity for now
          timestamp: Date.now(),
        });
      }
    }
    
    // Limit to max polyphony
    return detectedPitches.slice(0, this.options.maxPolyphony);
  }

  /**
   * Create spectrogram from time-domain signal
   */
  private createSpectrogram(samples: Float32Array): Float32Array[] {
    const hopSize = this.options.hopSize;
    const frameSize = this.options.frameSize;
    const spectrogram: Float32Array[] = [];
    
    // For single frame processing, create a simple spectrogram
    if (samples.length === frameSize) {
      // Apply window function
      const windowedFrame = this.applyWindow(samples);
      
      // Compute FFT magnitude
      const magnitude = this.computeFFTMagnitude(windowedFrame);
      
      spectrogram.push(magnitude);
    } else {
      // Extract overlapping frames for longer signals
      for (let i = 0; i <= samples.length - frameSize; i += hopSize) {
        const frame = samples.slice(i, i + frameSize);
        
        // Apply window function
        const windowedFrame = this.applyWindow(frame);
        
        // Compute FFT magnitude
        const magnitude = this.computeFFTMagnitude(windowedFrame);
        
        spectrogram.push(magnitude);
      }
    }
    
    return spectrogram;
  }

  /**
   * Apply window function to frame
   */
  private applyWindow(frame: Float32Array): Float32Array {
    const windowed = new Float32Array(frame.length);
    const window = this.createHannWindow(frame.length);
    
    for (let i = 0; i < frame.length; i++) {
      windowed[i] = frame[i] * window[i];
    }
    
    return windowed;
  }

  /**
   * Create Hann window
   */
  private createHannWindow(length: number): Float32Array {
    const window = new Float32Array(length);
    
    for (let i = 0; i < length; i++) {
      window[i] = 0.5 * (1 - Math.cos(2 * Math.PI * i / (length - 1)));
    }
    
    return window;
  }

  /**
   * Compute FFT magnitude (simplified implementation)
   */
  private computeFFTMagnitude(frame: Float32Array): Float32Array {
    // Simplified FFT implementation for testing
    // In a real implementation, this would use a proper FFT library
    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;
  }
}

/**
 * Utility function to convert frequency to MIDI note number
 */
export function frequencyToMidi(frequency: number): number {
  return Math.round(12 * Math.log2(frequency / 440) + 69);
}

/**
 * Utility function to convert MIDI note number to frequency
 */
export function midiToFrequency(midi: number): number {
  return 440 * Math.pow(2, (midi - 69) / 12);
}

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

// Re-export audio utilities
export { 
  AudioProcessor,
  mixToMono, 
  calculateRMS, 
  normalizeAudio,
  applyPreEmphasis,
  type AudioProcessorOptions 
} from './audio';

// Re-export worklet utilities
export {
  PitchWorklet,
  createPitchWorklet,
  type WorkletFrameEvent,
  type WorkletConfigMessage
} from './worklet';

// Re-export DSP utilities
export {
  FFT,
  WindowFunction,
  SpectralAnalysis,
  nextPowerOfTwo,
  isPowerOfTwo,
  type Complex,
  type WindowType,
  type FFTResult
} from './dsp';

// Re-export peak detection utilities
export {
  PeakDetector,
  findFundamental,
  type Peak,
  type PeakDetectionConfig
} from './peak-detection';

// Re-export TensorFlow.js utilities
export {
  TFJSModelManager,
  TensorUtils,
  createModelManager,
  tf,
  type ModelLoadOptions,
  type InferenceResult
} from './tfjs';

// Re-export CREPE utilities
export {
  CREPEModel,
  detectPitchCREPE,
  CREPEUtils,
  type CREPEConfig,
  type CREPEPrediction
} from './crepe';

// Re-export autocorrelation utilities
export {
  AutocorrelationDetector,
  detectPitchAutocorrelation,
  AutocorrelationUtils,
  type AutocorrelationConfig,
  type AutocorrelationResult
} from './autocorrelation';

// Re-export harmonic analysis utilities
export {
  HarmonicAnalyzer,
  analyzeHarmonics,
  HarmonicMatching,
  HarmonicUtils,
  type HarmonicAnalysisConfig,
  type HarmonicResult
} from './harmonic-analysis';

// Re-export pitch tracking utilities
export {
  PitchTracker,
  PitchTrackingUtils,
  trackPitch,
  type PitchTrackingConfig,
  type TrackedPitch,
  type PitchTrackingState
} from './pitch-tracking';

// Re-export NMF utilities
export {
  NMFAlgorithm,
  SpectralDictionary,
  NMFUtils,
  decomposeNMF,
  type NMFConfig,
  type NMFResult,
  type NMFComponent,
  type SpectralDictionary as SpectralDictionaryType
} from './nmf';

// Re-export CREPE-NMF utilities
export {
  CREPENMFDetector,
  detectPolyphonicPitches,
  type CREPENMFConfig,
  type CREPENMFResult,
  type CREPENMFPitch
} from './crepe-nmf';