/**
 * CREPE (Convolutional REcurrent Pitch Estimator) implementation
 * Neural network-based monophonic pitch detection
 */

import { TFJSModelManager, TensorUtils } from './tfjs';
import { WindowFunction } from './dsp';

/**
 * CREPE model configuration
 */
export interface CREPEConfig {
  modelUrl?: string;
  modelSize?: 'tiny' | 'small' | 'medium' | 'large' | 'full';
  viterbi?: boolean;
  centerFrequency?: number;
  stepSize?: number;
}

/**
 * CREPE prediction result
 */
export interface CREPEPrediction {
  frequency: number;
  confidence: number;
  activations?: Float32Array;
  time?: number;
}

/**
 * CREPE pitch bins configuration
 * 360 bins spanning 1951.2 Hz (20 cents per bin)
 */
const CREPE_BINS = 360;
const CREPE_MIN_FREQ = 32.70; // C1
const CREPE_MAX_FREQ = 1975.5; // ~B6
const CREPE_SAMPLE_RATE = 16000;
const CREPE_FRAME_SIZE = 1024;

/**
 * CREPE model wrapper
 */
export class CREPEModel {
  private modelManager: TFJSModelManager;
  private config: Required<CREPEConfig>;
  private isReady: boolean = false;

  constructor(config: CREPEConfig = {}) {
    this.config = {
      modelUrl: config.modelUrl ?? '/models/crepe/model.json',
      modelSize: config.modelSize ?? 'small',
      viterbi: config.viterbi ?? false,
      centerFrequency: config.centerFrequency ?? 0,
      stepSize: config.stepSize ?? 10, // ms
    };

    this.modelManager = new TFJSModelManager();
  }

  /**
   * Initialize and load CREPE model
   */
  async initialize(): Promise<void> {
    if (this.isReady) {
      return;
    }

    await this.modelManager.initialize('webgl'); // Prefer GPU for CREPE
    await this.modelManager.loadModel(this.config.modelUrl);
    
    this.isReady = true;
  }

  /**
   * Predict pitch from audio frame
   */
  async predict(audioFrame: Float32Array, sampleRate: number = 16000): Promise<CREPEPrediction> {
    if (!this.isReady) {
      throw new Error('CREPE model not initialized. Call initialize() first.');
    }

    // Resample to 16kHz if needed
    const resampled = this.resampleTo16kHz(audioFrame, sampleRate);

    // Extract 1024-sample frame with proper centering
    const frame = this.extractFrame(resampled);

    // Normalize frame
    const normalized = this.normalizeFrame(frame);

    // Run inference
    const activations = await this.modelManager.predictArray(normalized);

    // Convert activations to frequency and confidence
    const { frequency, confidence } = this.activationsToFrequency(activations);

    return {
      frequency,
      confidence,
      activations,
    };
  }

  /**
   * Predict pitch from multiple frames (batch processing)
   */
  async predictBatch(frames: Float32Array[], sampleRate: number = 16000): Promise<CREPEPrediction[]> {
    if (!this.isReady) {
      throw new Error('CREPE model not initialized. Call initialize() first.');
    }

    const results: CREPEPrediction[] = [];

    for (const frame of frames) {
      const result = await this.predict(frame, sampleRate);
      results.push(result);
    }

    return results;
  }

  /**
   * Resample audio to 16kHz (CREPE's expected sample rate)
   */
  private resampleTo16kHz(audio: Float32Array, sampleRate: number): Float32Array {
    if (sampleRate === CREPE_SAMPLE_RATE) {
      return audio;
    }

    const ratio = CREPE_SAMPLE_RATE / sampleRate;
    const outputLength = Math.floor(audio.length * ratio);
    const resampled = new Float32Array(outputLength);

    // Simple linear interpolation resampling
    for (let i = 0; i < outputLength; i++) {
      const srcIndex = i / ratio;
      const srcFloor = Math.floor(srcIndex);
      const srcCeil = Math.min(srcFloor + 1, audio.length - 1);
      const fraction = srcIndex - srcFloor;

      resampled[i] = audio[srcFloor] * (1 - fraction) + audio[srcCeil] * fraction;
    }

    return resampled;
  }

  /**
   * Extract 1024-sample frame from center of audio
   */
  private extractFrame(audio: Float32Array): Float32Array {
    const frame = new Float32Array(CREPE_FRAME_SIZE);

    if (audio.length >= CREPE_FRAME_SIZE) {
      // Center the frame
      const offset = Math.floor((audio.length - CREPE_FRAME_SIZE) / 2);
      frame.set(audio.slice(offset, offset + CREPE_FRAME_SIZE));
    } else {
      // Pad with zeros if too short
      frame.set(audio);
    }

    return frame;
  }

  /**
   * Normalize frame to [-1, 1] with mean removal
   */
  private normalizeFrame(frame: Float32Array): Float32Array {
    const normalized = new Float32Array(frame.length);

    // Remove mean
    let mean = 0;
    for (let i = 0; i < frame.length; i++) {
      mean += frame[i];
    }
    mean /= frame.length;

    // Normalize by max absolute value
    let maxAbs = 0;
    for (let i = 0; i < frame.length; i++) {
      const val = frame[i] - mean;
      normalized[i] = val;
      maxAbs = Math.max(maxAbs, Math.abs(val));
    }

    if (maxAbs > 0) {
      for (let i = 0; i < normalized.length; i++) {
        normalized[i] /= maxAbs;
      }
    }

    return normalized;
  }

  /**
   * Convert model activations to frequency and confidence
   */
  private activationsToFrequency(activations: Float32Array): { frequency: number; confidence: number } {
    // Find peak in activations
    let maxIdx = 0;
    let maxVal = activations[0];

    for (let i = 1; i < activations.length; i++) {
      if (activations[i] > maxVal) {
        maxVal = activations[i];
        maxIdx = i;
      }
    }

    // Weighted average around peak for sub-bin accuracy
    const frequency = this.binToFrequency(maxIdx, activations);
    const confidence = maxVal;

    return { frequency, confidence };
  }

  /**
   * Convert bin index to frequency with interpolation
   */
  private binToFrequency(bin: number, activations: Float32Array): number {
    // Use parabolic interpolation for better accuracy
    if (bin > 0 && bin < activations.length - 1) {
      const alpha = activations[bin - 1];
      const beta = activations[bin];
      const gamma = activations[bin + 1];

      const p = 0.5 * (alpha - gamma) / (alpha - 2 * beta + gamma);
      const interpolatedBin = bin + p;

      // Convert bin to frequency (CREPE uses 20 cents per bin)
      const centsPerBin = 20;
      const cents = interpolatedBin * centsPerBin;
      return CREPE_MIN_FREQ * Math.pow(2, cents / 1200);
    }

    // No interpolation for edge bins
    const centsPerBin = 20;
    const cents = bin * centsPerBin;
    return CREPE_MIN_FREQ * Math.pow(2, cents / 1200);
  }

  /**
   * Convert frequency to CREPE bin
   */
  frequencyToBin(frequency: number): number {
    const cents = 1200 * Math.log2(frequency / CREPE_MIN_FREQ);
    return Math.round(cents / 20);
  }

  /**
   * Apply Viterbi smoothing to sequence of predictions
   */
  viterbiSmooth(predictions: CREPEPrediction[]): CREPEPrediction[] {
    if (!this.config.viterbi || predictions.length < 2) {
      return predictions;
    }

    // Simple Viterbi-like smoothing
    const smoothed: CREPEPrediction[] = [];
    let prevFreq = predictions[0].frequency;

    for (const pred of predictions) {
      // Penalize large jumps in frequency
      const jump = Math.abs(pred.frequency - prevFreq);
      const penalty = jump / prevFreq;
      const adjustedConfidence = pred.confidence * Math.exp(-penalty);

      smoothed.push({
        ...pred,
        confidence: adjustedConfidence,
      });

      prevFreq = pred.frequency;
    }

    return smoothed;
  }

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

  /**
   * Check if model is ready
   */
  isInitialized(): boolean {
    return this.isReady;
  }

  /**
   * Dispose model and free memory
   */
  dispose(): void {
    this.modelManager.dispose();
    this.isReady = false;
  }
}

/**
 * Quick utility: Get pitch from audio using CREPE
 */
export async function detectPitchCREPE(
  audio: Float32Array,
  sampleRate: number = 16000,
  modelUrl?: string
): Promise<number | null> {
  const crepe = new CREPEModel({ modelUrl });
  await crepe.initialize();

  const result = await crepe.predict(audio, sampleRate);
  crepe.dispose();

  return result.confidence > 0.5 ? result.frequency : null;
}

/**
 * CREPE frequency range utilities
 */
export const CREPEUtils = {
  MIN_FREQ: CREPE_MIN_FREQ,
  MAX_FREQ: CREPE_MAX_FREQ,
  SAMPLE_RATE: CREPE_SAMPLE_RATE,
  FRAME_SIZE: CREPE_FRAME_SIZE,
  BINS: CREPE_BINS,

  /**
   * Check if frequency is in CREPE range
   */
  isInRange(frequency: number): boolean {
    return frequency >= CREPE_MIN_FREQ && frequency <= CREPE_MAX_FREQ;
  },

  /**
   * Get nearest valid CREPE frequency
   */
  clamp(frequency: number): number {
    return Math.max(CREPE_MIN_FREQ, Math.min(CREPE_MAX_FREQ, frequency));
  },

  /**
   * Convert MIDI to CREPE bin
   */
  midiToBin(midi: number): number {
    const frequency = 440 * Math.pow(2, (midi - 69) / 12);
    const cents = 1200 * Math.log2(frequency / CREPE_MIN_FREQ);
    return Math.round(cents / 20);
  },

  /**
   * Convert CREPE bin to MIDI
   */
  binToMidi(bin: number): number {
    const cents = bin * 20;
    const frequency = CREPE_MIN_FREQ * Math.pow(2, cents / 1200);
    return Math.round(12 * Math.log2(frequency / 440) + 69);
  },
};