// MQ Extraction Algorithm
// McAulay-Quatieri sinusoidal analysis

// Extract partials from audio buffer
function extractPartials(params, stftCache) {
  const {fftSize, threshold, sampleRate, freqWeight, prominence} = params;
  const numFrames = stftCache.getNumFrames();

  const frames = [];
  for (let i = 0; i < numFrames; ++i) {
    const cachedFrame = stftCache.getFrameAtIndex(i);
    const squaredAmp = cachedFrame.squaredAmplitude;
    const phase = cachedFrame.phase;
    const peaks = detectPeaks(squaredAmp, phase, fftSize, sampleRate, threshold, freqWeight, prominence);
    frames.push({time: cachedFrame.time, peaks});
  }

  const partials = trackPartials(frames, params);

  // Second pass: extend partials leftward to recover onset frames
  expandPartialsLeft(partials, frames);

  for (const partial of partials) {
    partial.freqCurve = fitBezier(partial.times, partial.freqs);
    partial.ampCurve = fitBezier(partial.times, partial.amps);
  }

  return {partials, frames};
}

// Helper to interpolate phase via quadratic formula on unwrapped neighbors.
// This provides a more accurate phase estimate at the sub-bin peak location.
function phaseInterp(p_minus, p_center, p_plus, p_frac) {
  // unwrap neighbors relative to center
  let dp_minus = p_minus - p_center;
  while (dp_minus > Math.PI) dp_minus -= 2 * Math.PI;
  while (dp_minus < -Math.PI) dp_minus += 2 * Math.PI;

  let dp_plus = p_plus - p_center;
  while (dp_plus > Math.PI) dp_plus -= 2 * Math.PI;
  while (dp_plus < -Math.PI) dp_plus += 2 * Math.PI;
  
  const p_interp = p_center + (dp_plus - dp_minus) * p_frac * 0.5 + (dp_plus + dp_minus) * p_frac * p_frac;
  return p_interp;
}

// Detect spectral peaks via local maxima + parabolic interpolation
// squaredAmp: pre-computed re*re+im*im per bin
// phase: pre-computed atan2(im,re) per bin
// freqWeight: if true, weight by f before peak detection (f * Power(f))
function detectPeaks(squaredAmp, phase, fftSize, sampleRate, thresholdDB, freqWeight, prominenceDB = 0) {
  const mag = new Float32Array(fftSize / 2);
  const binHz = sampleRate / fftSize;
  for (let i = 0; i < fftSize / 2; ++i) {
    const w = freqWeight ? (i * binHz) : 1.0;
    mag[i] = 10 * Math.log10(Math.max(squaredAmp[i] * w, 1e-20));
  }

  const peaks = [];
  for (let i = 2; i < mag.length - 2; ++i) {
    if (mag[i] > thresholdDB &&
        mag[i] > mag[i-1] && mag[i] > mag[i-2] &&
        mag[i] > mag[i+1] && mag[i] > mag[i+2]) {

      // Check prominence if requested
      if (prominenceDB > 0) {
        let minLeft = mag[i];
        for (let k = i - 1; k >= 0; --k) {
          if (mag[k] > mag[i]) break; // Found higher peak
          if (mag[k] < minLeft) minLeft = mag[k];
        }

        let minRight = mag[i];
        for (let k = i + 1; k < mag.length; ++k) {
          if (mag[k] > mag[i]) break; // Found higher peak
          if (mag[k] < minRight) minRight = mag[k];
        }

        const valley = Math.max(minLeft, minRight);
        if (mag[i] - valley < prominenceDB) continue;
      }

      // Parabolic interpolation for sub-bin accuracy on frequency, amplitude, and phase
      const alpha = mag[i-1];
      const beta  = mag[i];
      const gamma = mag[i+1];
      const p = 0.5 * (alpha - gamma) / (alpha - 2*beta + gamma);
      
      const p_phase = phaseInterp(phase[i-1], phase[i], phase[i+1], p);

      const freq  = (i + p) * sampleRate / fftSize;
      const ampDB = beta - 0.25 * (alpha - gamma) * p;
      peaks.push({freq, amp: Math.pow(10, ampDB / 20), phase: p_phase});
    }
  }

  return peaks;
}

// Helper to compute shortest angle difference (e.g., between -pi and pi)
function normalizeAngle(angle) {
  return angle - 2 * Math.PI * Math.floor((angle + Math.PI) / (2 * Math.PI));
}

// Track partials across frames using phase coherence for robust matching.
function trackPartials(frames, params) {
  const { sampleRate, hopSize } = params;
  const partials        = [];
  const activePartials  = [];
  const candidates      = []; // pre-birth

  const trackingRatio   = 0.05; // 5% frequency tolerance
  const minTrackingHz   = 20;
  const birthPersistence = 3;   // frames before partial is born
  const deathAge        = 5;    // frames without match before death
  const minLength       = 10;   // frames required to keep partial

  // Weight phase error heavily in cost function, scaled by frequency.
  // This makes phase deviation more significant for high-frequency partials.
  const phaseErrorWeight = 2.0;

  for (const frame of frames) {
    const matched = new Set();

    // --- Continue active partials ---
    for (const partial of activePartials) {
      const lastFreq = partial.freqs[partial.freqs.length - 1];
      const lastPhase = partial.phases[partial.phases.length - 1];
      const velocity = partial.velocity || 0;
      const predictedFreq = lastFreq + velocity;
      
      // Predict phase for the current frame based on the last frame's frequency.
      const phaseAdvance = 2 * Math.PI * lastFreq * hopSize / sampleRate;
      const predictedPhase = lastPhase + phaseAdvance;
      
      const tol = Math.max(predictedFreq * trackingRatio, minTrackingHz);
      let bestIdx = -1, bestCost = Infinity;

      // Find the peak in the new frame with the lowest cost (freq + phase error).
      for (let i = 0; i < frame.peaks.length; ++i) {
        if (matched.has(i)) continue;
        const pk = frame.peaks[i];
        const freqError = Math.abs(pk.freq - predictedFreq);
        if (freqError > tol) continue;

        const phaseError = Math.abs(normalizeAngle(pk.phase - predictedPhase));
        const cost = freqError + phaseErrorWeight * phaseError * predictedFreq;
        if (cost < bestCost) { bestCost = cost; bestIdx = i; }
      }

      if (bestIdx >= 0) {
        const pk = frame.peaks[bestIdx];
        partial.times.push(frame.time);
        partial.freqs.push(pk.freq);
        partial.amps.push(pk.amp);
        partial.phases.push(pk.phase);
        partial.age = 0;
        partial.velocity = pk.freq - lastFreq;
        matched.add(bestIdx);
      } else {
        partial.age++;
      }
    }

    // --- Advance candidates ---
    for (let i = candidates.length - 1; i >= 0; --i) {
      const cand = candidates[i];
      const lastFreq = cand.freqs[cand.freqs.length - 1];
      const lastPhase = cand.phases[cand.phases.length - 1];
      const velocity = cand.velocity || 0;
      const predictedFreq = lastFreq + velocity;

      const phaseAdvance = 2 * Math.PI * lastFreq * hopSize / sampleRate;
      const predictedPhase = lastPhase + phaseAdvance;

      const tol = Math.max(predictedFreq * trackingRatio, minTrackingHz);
      let bestIdx = -1, bestCost = Infinity;

      for (let j = 0; j < frame.peaks.length; ++j) {
        if (matched.has(j)) continue;
        const pk = frame.peaks[j];
        const freqError = Math.abs(pk.freq - predictedFreq);
        if (freqError > tol) continue;
        
        const phaseError = Math.abs(normalizeAngle(pk.phase - predictedPhase));
        const cost = freqError + phaseErrorWeight * phaseError * predictedFreq;
        if (cost < bestCost) { bestCost = cost; bestIdx = j; }
      }

      if (bestIdx >= 0) {
        const pk = frame.peaks[bestIdx];
        cand.times.push(frame.time);
        cand.freqs.push(pk.freq);
        cand.amps.push(pk.amp);
        cand.phases.push(pk.phase);
        cand.velocity = pk.freq - lastFreq;
        matched.add(bestIdx);
        // "graduate" a candidate to a full partial
        if (cand.times.length >= birthPersistence) {
          activePartials.push(cand);
          candidates.splice(i, 1);
        }
      } else {
        candidates.splice(i, 1); // kill candidate
      }
    }

    // --- Spawn new candidates from unmatched peaks ---
    for (let i = 0; i < frame.peaks.length; ++i) {
      if (matched.has(i)) continue;
      const pk = frame.peaks[i];
      candidates.push({
        times: [frame.time], 
        freqs: [pk.freq], 
        amps: [pk.amp],
        phases: [pk.phase],
        age: 0,
        velocity: 0
      });
    }

    // --- Kill aged-out partials ---
    for (let i = activePartials.length - 1; i >= 0; --i) {
      if (activePartials[i].age > deathAge) {
        if (activePartials[i].times.length >= minLength) partials.push(activePartials[i]);
        activePartials.splice(i, 1);
      }
    }
  }

  // --- Collect remaining active partials ---
  for (const partial of activePartials) {
    if (partial.times.length >= minLength) partials.push(partial);
  }

  return partials;
}

// Second pass: extend each partial leftward to recover onset frames missed
// by the birthPersistence requirement in the forward pass.
function expandPartialsLeft(partials, frames) {
  const trackingRatio = 0.05;
  const minTrackingHz = 20;

  // Build time → frame index map
  const timeToIdx = new Map();
  for (let i = 0; i < frames.length; ++i) timeToIdx.set(frames[i].time, i);

  for (const partial of partials) {
    if (!partial.phases) partial.phases = []; // Ensure old partials have phase array

    let startIdx = timeToIdx.get(partial.times[0]);
    if (startIdx == null || startIdx === 0) continue;

    for (let i = startIdx - 1; i >= 0; --i) {
      const frame = frames[i];
      const refFreq = partial.freqs[0];
      const tol = Math.max(refFreq * trackingRatio, minTrackingHz);

      let bestIdx = -1, bestDist = Infinity;
      for (let j = 0; j < frame.peaks.length; ++j) {
        const dist = Math.abs(frame.peaks[j].freq - refFreq);
        if (dist < tol && dist < bestDist) { bestDist = dist; bestIdx = j; }
      }

      if (bestIdx < 0) break;

      const pk = frame.peaks[bestIdx];
      partial.times.unshift(frame.time);
      partial.freqs.unshift(pk.freq);
      partial.amps.unshift(pk.amp);
      partial.phases.unshift(pk.phase);
    }
  }
}

// Autodetect spread_above / spread_below from the spectrogram.
// For each (subsampled) STFT frame within the partial, measures the
// half-power (-3dB) width of the spectral peak above and below the center.
// spread = half_bandwidth / f0  (fractional).
function autodetectSpread(partial, stftCache, fftSize, sampleRate) {
  const curve = partial.freqCurve;
  if (!curve || !stftCache) return {spread_above: 0.02, spread_below: 0.02};

  const numFrames = stftCache.getNumFrames();
  const binHz    = sampleRate / fftSize;
  const halfBins = fftSize / 2;

  let sumAbove = 0, sumBelow = 0, count = 0;

  const STEP = 4;
  for (let fi = 0; fi < numFrames; fi += STEP) {
    const frame = stftCache.getFrameAtIndex(fi);
    if (!frame) continue;
    const t = frame.time;
    if (t < curve.t0 || t > curve.t3) continue;

    const f0 = evalBezier(curve, t);
    if (f0 <= 0) continue;

    const sq = frame.squaredAmplitude;
    if (!sq) continue;

    // Find peak bin in ±10% window
    const binCenter = f0 / binHz;
    const searchBins = Math.max(3, Math.round(f0 * 0.10 / binHz));
    const binLo = Math.max(1, Math.floor(binCenter - searchBins));
    const binHi = Math.min(halfBins - 2, Math.ceil(binCenter + searchBins));

    let peakBin = binLo, peakVal = sq[binLo];
    for (let b = binLo + 1; b <= binHi; ++b) {
      if (sq[b] > peakVal) { peakVal = sq[b]; peakBin = b; }
    }

    const halfPower = peakVal * 0.5;  // -3dB in power

    // Walk above peak until half-power, interpolate crossing
    let aboveBin = peakBin;
    while (aboveBin < halfBins - 1 && sq[aboveBin] > halfPower) ++aboveBin;
    const tA = aboveBin > peakBin && sq[aboveBin - 1] !== sq[aboveBin]
      ? (halfPower - sq[aboveBin - 1]) / (sq[aboveBin] - sq[aboveBin - 1])
      : 0;
    const widthAbove = (aboveBin - 1 + tA - peakBin) * binHz;

    // Walk below peak until half-power, interpolate crossing
    let belowBin = peakBin;
    while (belowBin > 1 && sq[belowBin] > halfPower) --belowBin;
    const tB = belowBin < peakBin && sq[belowBin + 1] !== sq[belowBin]
      ? (halfPower - sq[belowBin + 1]) / (sq[belowBin] - sq[belowBin + 1])
      : 0;
    const widthBelow = (peakBin - belowBin - 1 + tB) * binHz;

    sumAbove += (widthAbove / f0) * (widthAbove / f0);
    sumBelow += (widthBelow / f0) * (widthBelow / f0);
    ++count;
  }

  const spread_above = count > 0 ? Math.sqrt(sumAbove / count) : 0.01;
  const spread_below = count > 0 ? Math.sqrt(sumBelow / count) : 0.01;
  return {spread_above, spread_below};
}

// Fit cubic bezier to trajectory using least-squares for inner control points
function fitBezier(times, values) {
  const n = times.length - 1;
  const t0 = times[0], v0 = values[0];
  const t3 = times[n], v3 = values[n];
  const dt = t3 - t0;

  if (dt <= 1e-9 || n < 2) {
    // Linear fallback for too few points or zero duration
    return {t0, v0, t1: t0 + dt / 3, v1: v0 + (v3 - v0) / 3, t2: t0 + 2 * dt / 3, v2: v0 + 2 * (v3 - v0) / 3, t3, v3};
  }

  // Least squares solve for v1, v2
  // Bezier: B(u) = (1-u)^3*v0 + 3(1-u)^2*u*v1 + 3(1-u)*u^2*v2 + u^3*v3
  // Target_i = val_i - (1-u)^3*v0 - u^3*v3
  // Model_i = A_i*v1 + B_i*v2
  // A_i = 3(1-u)^2*u
  // B_i = 3(1-u)*u^2

  let sA2 = 0, sB2 = 0, sAB = 0, sAT = 0, sBT = 0;

  for (let i = 0; i <= n; ++i) {
    const u = (times[i] - t0) / dt;
    const u2 = u * u;
    const u3 = u2 * u;
    const invU = 1.0 - u;
    const invU2 = invU * invU;
    const invU3 = invU2 * invU;

    const A = 3 * invU2 * u;
    const B = 3 * invU * u2;
    const target = values[i] - (invU3 * v0 + u3 * v3);

    sA2 += A * A;
    sB2 += B * B;
    sAB += A * B;
    sAT += A * target;
    sBT += B * target;
  }

  const det = sA2 * sB2 - sAB * sAB;
  let v1, v2;

  if (Math.abs(det) < 1e-9) {
    // Fallback to simple 1/3, 2/3 heuristic if matrix is singular
    const idx1 = Math.round(n / 3);
    const idx2 = Math.round(2 * n / 3);
    v1 = values[idx1];
    v2 = values[idx2];
  } else {
    v1 = (sB2 * sAT - sAB * sBT) / det;
    v2 = (sA2 * sBT - sAB * sAT) / det;
  }

  return {t0, v0, t1: t0 + dt / 3, v1, t2: t0 + 2 * dt / 3, v2, t3, v3};
}