// YIN pitch detection algorithm

// Based on: de Cheveigné & Kawahara (2002)

// "YIN, a fundamental frequency estimator for speech and music"

#include "autocorrelation.hh"

#include "narrow.hh"

#include "xrange.hh"

#include <algorithm>

#include <array>

#include <bit>

#include <cassert>

#include <cmath>

#include <span>

namespace yin {

template<typename T> [[nodiscard]] constexpr T square(T t)

{

	return t * t;

}

// Helper: Compute difference function d(τ) using autocorrelation

//

// d(τ) = Σ(x[n] - x[n+τ])²

//      = Σx[n]² + Σx[n+τ]² - 2·Σx[n]·x[n+τ]

//      = energyLeft(τ) + energyRight(τ) - 2·r(τ)

//

// Returns: Total energy of the signal (sum of all squared samples)

//

// Preconditions:

//   - autocorr.size() == diff.size()

//   - diff.size() <= signal.size()

//   - WORK_LEN >= signal.size() (for internal cumSumSq buffer)

template<size_t WORK_LEN>

[[nodiscard]] inline float computeDifferenceFunction(

	std::span<const float> signal,

	std::span<const float> autocorr,

	std::span<float> diff)

{

	assert(autocorr.size() == diff.size());

	assert(diff.size() <= signal.size());

	const auto N = signal.size();

	const auto maxLag = diff.size();

	assert(WORK_LEN >= N);

	// Internal work buffer for cumulative sum of squares

	std::array<float, WORK_LEN> cumSumSq; // ok, uninitialized

	// Precompute cumulative sum of squares

	// cumSumSq[i] = Σ signal[j]² for j=0..i (inclusive)

	cumSumSq[0] = square(signal[0]);

	for (auto i : xrange(size_t(1), N)) {

		cumSumSq[i] = cumSumSq[i - 1] + square(signal[i]);

	}

	const auto totalEnergy = cumSumSq[N - 1];

	// Compute difference function

	diff[0] = 0.0f; // Perfect match at lag 0

	for (auto lag : xrange(size_t(1), maxLag)) {

		// energyLeft = Σ signal[n]² for n in [0, N-lag)

		auto energyLeft = cumSumSq[N - lag - 1];

		// energyRight = Σ signal[n]² for n in [lag, N)

		auto energyRight = totalEnergy - cumSumSq[lag - 1];

		// d(τ) = energyLeft + energyRight - 2·r(τ)

		diff[lag] = energyLeft + energyRight - 2.0f * autocorr[lag];

	}

	return totalEnergy;

}

// Helper: Cumulative mean normalization

//

// d'(τ) = d(τ) / [(1/τ) · Σd(j) for j=1..τ]

//       = d(τ) · τ / Σd(j) for j=1..τ

//

// This is the key innovation of YIN that reduces octave errors.

inline void cumulativeMeanNormalize(std::span<const float> diff,

                                    std::span<float> diffNorm)

{

	assert(diff.size() == diffNorm.size());

	diffNorm[0] = 1.0f; // Special case

	auto runningSum = 0.0f;

	for (auto lag : xrange(size_t(1), diff.size())) {

		runningSum += diff[lag];

		if (runningSum > 0.0f) {

			diffNorm[lag] = diff[lag] * float(lag) / runningSum;

		} else {

			diffNorm[lag] = 1.0f;

		}

	}

}

// Helper: Find first minimum below threshold

//

// Returns the lag corresponding to the fundamental period.

// Uses YIN's absolute threshold approach.

// Always returns a result (the best candidate found, even if weak).

//

// Preconditions:

//   - minLag < maxLag

//   - diffNorm.size() >= maxLag (typically size == maxLag + 1)

[[nodiscard]] inline size_t findBestLag(

	std::span<const float> diffNorm, size_t minLag,

	size_t maxLag, float threshold)

{

	assert(minLag < maxLag);

	assert(diffNorm.size() >= maxLag);

	// Single-pass: look for first value below threshold (with early exit)

	// while tracking global minimum as fallback

	auto bestLag = minLag;

	auto minVal = diffNorm[minLag];

	for (auto lag : xrange(minLag, maxLag)) {

		auto val = diffNorm[lag];

		// Check for first value below threshold -> early exit

		if (val < threshold) {

			// Found candidate, search forward for local minimum

			while (lag + 1 < maxLag &&

			       diffNorm[lag + 1] < diffNorm[lag]) {

				++lag;

			}

			return lag; // Early exit!

		}

		// Track global minimum (used if nothing below threshold)

		if (val < minVal) {

			minVal = val;

			bestLag = lag;

		}

	}

	// No value below threshold, use tracked global minimum

	// (Return even if confidence is low - let caller decide if acceptable)

	return bestLag;

}

// Helper: Parabolic interpolation for sub-sample accuracy

//

// Uses 3 points around the minimum to fit a parabola and find the true minimum.

[[nodiscard]] inline float parabolicInterpolation(

	std::span<const float> diffNorm, size_t lag)

{

	// Need neighbors on both sides (check array bounds only)

	if (lag == 0 || lag >= diffNorm.size() - 1) {

		return float(lag);

	}

	auto y1 = diffNorm[lag - 1];

	auto y2 = diffNorm[lag + 0];

	auto y3 = diffNorm[lag + 1];

	// Check we have a proper minimum (not maximum or flat)

	auto denom = y1 - 2.0f * y2 + y3;

	if (std::abs(denom) <= 1e-10f || y1 <= y2 || y3 <= y2) {

		// Can't fit parabola or not a minimum

		return float(lag);

	}

	// Parabolic interpolation formula

	// The offset is mathematically guaranteed to be in (-1, 1) given the

	// guards above: For a concave-up parabola where the middle of three

	// equally-spaced points is strictly lower than both neighbors, the

	// vertex must lie between them.

	auto offset = 0.5f * (y1 - y3) / denom;

	return float(lag) + offset;

}

struct PitchResult {

	float frequency = 0.0f; // Detected frequency in Hz (0.0 = no pitch detected)

	float error = 1.0f; // YIN aperiodicity measure (0.0 = perfect periodic,

	                    // higher = worse) Typical: <0.1 excellent, 0.1-0.2

	                    // good, 0.2-0.5 weak, >0.5 poor Default 1.0 = no

	                    // valid pitch detected

};

// Template helper that does the actual YIN computation for a specific FFT size

template<unsigned FFT_LEN_L2>

[[nodiscard]] inline PitchResult detectPitchImpl(

	std::span<const float> signal, float sampleRate,

	size_t minLag, size_t maxLag, float threshold)

{

	static constexpr unsigned FFT_LEN = 1 << FFT_LEN_L2;

	static constexpr unsigned WORK_LEN = FFT_LEN / 2;

	// FFT requires full FFT_LEN, but YIN working buffers only need WORK_LEN

	// because:

	//   - FFT_LEN >= 2*N (chosen to be next power of 2 above 2*N)

	//   - So N <= FFT_LEN/2, and maxLag <= N/2 <= FFT_LEN/4

	//   - Therefore maxLag + 1 < WORK_LEN, and signal.size() <= WORK_LEN

	// Compute autocorrelation using the new autocorrelation<SIGNAL_LEN_L2> API

	// SIGNAL_LEN_L2 = FFT_LEN_L2 - 1 because FFT_LEN = 2 * SIGNAL_LEN

	std::array<float, WORK_LEN> autocorr_; // ok, uninitialized

	auto autocorrFull = std::span(autocorr_).first(signal.size());

	autocorrelation<FFT_LEN_L2 - 1>(signal, autocorrFull);

	// Extract only the lags we need (0..maxLag)

	auto autocorr = autocorrFull.first(maxLag + 1);

	// Compute difference function

	std::array<float, WORK_LEN> diff_; // ok, uninitialized

	auto diff = std::span(diff_).first(maxLag + 1);

	auto totalEnergy = computeDifferenceFunction<WORK_LEN>(signal, autocorr, diff);

	// Early exit if signal is essentially silence

	if (totalEnergy < 1e-10f) return {};

	// Cumulative mean normalization

	std::array<float, WORK_LEN> diffNorm_; // ok, uninitialized

	auto diffNorm = std::span(diffNorm_).first(maxLag + 1);

	cumulativeMeanNormalize(diff, diffNorm);

	// Find best lag (always returns a result, even if confidence is low)

	auto bestLag = findBestLag(diffNorm, minLag, maxLag, threshold);

	// Parabolic interpolation

	auto refinedLag = parabolicInterpolation(diffNorm, bestLag);

	// Convert to frequency and get error measure

	auto frequency = sampleRate / refinedLag;

	auto error = std::max(diffNorm[bestLag], 0.0f); // YIN normalized difference (lower is better)

	return {frequency, error};

}

// Main YIN pitch detection function (public API)

//

// Detects the fundamental frequency of an audio signal using the YIN algorithm.

//

// Parameters:

//   signal     - Audio samples

//   sampleRate - Sample rate in Hz (e.g., 44100.0f)

//   minFreq    - Minimum frequency to search for (default: 50 Hz)

//   maxFreq    - Maximum frequency to search for (default: 2000 Hz)

//   threshold  - YIN threshold parameter (default: 0.15, range: 0.1-0.2)

//

// Returns:

//   PitchResult with:

//     frequency  - Detected pitch in Hz

//     error      - YIN aperiodicity measure (lower is better)

//                  Typical values: <0.1 = excellent, 0.1-0.2 = good, 0.2-0.5 =

//                  weak, >0.5 = poor Caller should filter out results with high

//                  error (e.g., >0.5) Can exceed 1.0 for very non-periodic

//                  signals (drums, noise)

//

// Preconditions (checked by assertions in debug builds):

//   - signal.size() >= 2

//   - minFreq < maxFreq

//   - signal.size() >= 2 * sampleRate / minFreq (approximately)

//     i.e., signal must contain at least ~2 periods of the lowest frequency

//   - Recommended buffer size: ~2048 samples @ 44.1kHz for good frequency

//   resolution

//     (allows detecting down to ~43 Hz with 2048 samples @ 44.1kHz)

[[nodiscard]] inline PitchResult detectPitch(

	std::span<const float> signal,

	float sampleRate,

	float minFreq = 50.0f,

	float maxFreq = 2000.0f,

	float threshold = 0.15f)

{

	const auto N = signal.size();

	// Calculate lag range

	const auto minLag = std::max(size_t(1), size_t(sampleRate / maxFreq));

	const auto maxLag = std::min(size_t(sampleRate / minFreq), N / 2);

	// Preconditions: signal must be large enough and frequency range must

	// be valid

	assert(N >= 2);

	assert(maxLag > minLag);

	// Determine FFT size: next power of 2 that's at least 2*N

	const auto fftLenL2 = narrow<unsigned>(std::bit_width(2 * N - 1));

	// Dispatch to appropriate FFT size

	switch (fftLenL2) {

	case  8: return detectPitchImpl< 8>(signal, sampleRate, minLag, maxLag, threshold);

	case  9: return detectPitchImpl< 9>(signal, sampleRate, minLag, maxLag, threshold);

	case 10: return detectPitchImpl<10>(signal, sampleRate, minLag, maxLag, threshold);

	case 11: return detectPitchImpl<11>(signal, sampleRate, minLag, maxLag, threshold);

	case 12: return detectPitchImpl<12>(signal, sampleRate, minLag, maxLag, threshold);

	case 13: return detectPitchImpl<13>(signal, sampleRate, minLag, maxLag, threshold);

	case 14: return detectPitchImpl<14>(signal, sampleRate, minLag, maxLag, threshold);

	default:

		// Unsupported FFT size - signal too large or too small

		assert(false && "Unsupported FFT size in detectPitchYIN");

		return {};

	}

}

} // namespace yin