rosa 0.1.0

//! YIN and pYIN pitch estimation matching `librosa.yin` and `librosa.pyin`.
//!
//! Both algorithms share internal helpers (CMND, parabolic interpolation, trough
//! finding), so they live in the same module to avoid duplication.

use crate::matrix::Matrix;

/// Parameters for YIN pitch estimation.
pub struct YinParams {
    /// Sampling rate in Hz. Default: 22050.0
    pub sr: f64,
    /// Minimum frequency in Hz.
    pub fmin: f64,
    /// Maximum frequency in Hz.
    pub fmax: f64,
    /// Frame length in samples. Default: 2048
    pub frame_length: usize,
    /// Hop length in samples. Default: frame_length / 4
    pub hop_length: Option<usize>,
    /// Absolute threshold for trough selection. Default: 0.1
    pub trough_threshold: f64,
    /// Center-pad the signal. Default: true
    pub center: bool,
}

impl YinParams {
    /// Create a new `YinParams` with the given frequency range and default values.
    pub fn new(fmin: f64, fmax: f64) -> Self {
        Self {
            sr: 22050.0,
            fmin,
            fmax,
            frame_length: 2048,
            hop_length: None,
            trough_threshold: 0.1,
            center: true,
        }
    }

    fn hop(&self) -> usize {
        self.hop_length.unwrap_or(self.frame_length / 4)
    }
}

/// Parameters for pYIN pitch estimation.
pub struct PyinParams {
    /// Sampling rate in Hz. Default: 22050.0
    pub sr: f64,
    /// Minimum frequency in Hz.
    pub fmin: f64,
    /// Maximum frequency in Hz.
    pub fmax: f64,
    /// Frame length in samples. Default: 2048
    pub frame_length: usize,
    /// Hop length in samples. Default: frame_length / 4
    pub hop_length: Option<usize>,
    /// Number of thresholds for candidate generation. Default: 100
    pub n_thresholds: usize,
    /// Beta distribution shape parameters. Default: (2.0, 18.0)
    pub beta_parameters: (f64, f64),
    /// Boltzmann distribution parameter. Default: 2.0
    pub boltzmann_parameter: f64,
    /// Pitch bin resolution (fraction of semitone). Default: 0.1
    pub resolution: f64,
    /// Maximum pitch transition rate in octaves/second. Default: 35.92
    pub max_transition_rate: f64,
    /// Voiced/unvoiced switching probability. Default: 0.01
    pub switch_prob: f64,
    /// Probability for no-trough fallback. Default: 0.01
    pub no_trough_prob: f64,
    /// Value for unvoiced frames. Default: Some(NaN)
    pub fill_na: Option<f64>,
    /// Center-pad the signal. Default: true
    pub center: bool,
}

impl PyinParams {
    /// Create a new `PyinParams` with the given frequency range and default values.
    pub fn new(fmin: f64, fmax: f64) -> Self {
        Self {
            sr: 22050.0,
            fmin,
            fmax,
            frame_length: 2048,
            hop_length: None,
            n_thresholds: 100,
            beta_parameters: (2.0, 18.0),
            boltzmann_parameter: 2.0,
            resolution: 0.1,
            max_transition_rate: 35.92,
            switch_prob: 0.01,
            no_trough_prob: 0.01,
            fill_na: Some(f64::NAN),
            center: true,
        }
    }

    fn hop(&self) -> usize {
        self.hop_length.unwrap_or(self.frame_length / 4)
    }
}

/// Result of pYIN pitch estimation.
pub struct PyinResult {
    /// Fundamental frequency per frame (Hz). NaN for unvoiced frames.
    pub f0: Vec<f64>,
    /// Whether each frame is voiced.
    pub voiced_flag: Vec<bool>,
    /// Probability of voicing per frame.
    pub voiced_prob: Vec<f64>,
}

// --- Internal helpers ---

/// Parabolic interpolation: compute sub-sample shift for each element.
///
/// For each element k (excluding boundaries), fits a parabola through
/// (k-1, k, k+1) and returns the offset of the vertex from k.
///
/// The guard `|b| < |a|` ensures the parabola opens away from the vertex
/// (concave for minima) and that the shift stays within ±0.5 samples.
/// When `|b| >= |a|`, the fit is unreliable (nearly flat or wrong curvature)
/// so shift defaults to 0.
fn parabolic_interpolation(x: &[f64]) -> Vec<f64> {
    let n = x.len();
    let mut shifts = vec![0.0; n];
    for k in 1..n.saturating_sub(1) {
        let a = x[k - 1] + x[k + 1] - 2.0 * x[k];
        let b = (x[k + 1] - x[k - 1]) / 2.0;
        if b.abs() < a.abs() {
            shifts[k] = -b / a;
        }
    }
    shifts
}

/// Cumulative Mean Normalized Difference function (YIN core).
///
/// Returns CMND values for periods [min_period..=max_period].
fn cumulative_mean_normalized_difference(
    frame: &[f64],
    min_period: usize,
    max_period: usize,
) -> Vec<f64> {
    // Autocorrelate with max_size = max_period + 1
    let acf = crate::autocorrelate::autocorrelate(frame, Some(max_period + 1));

    // Compute cumulative energy of the frame
    let energy_sq: Vec<f64> = frame.iter().map(|x| x * x).collect();
    let mut cumsum_sq = vec![0.0; energy_sq.len()];
    let mut running = 0.0;
    for (i, &v) in energy_sq.iter().enumerate() {
        running += v;
        cumsum_sq[i] = running;
    }

    // Compute difference function d(tau) for tau = 1..max_period
    // d(tau) = 2*(acf[0] - acf[tau])
    // But librosa's implementation also subtracts a cumulative energy term.
    // d[0] = 0
    let n_tau = max_period + 1;
    let mut yin_frame = vec![0.0; n_tau];
    yin_frame[0] = 0.0;

    for tau in 1..n_tau {
        // d(tau) = 2*(ACF(0) - ACF(tau))
        let diff = 2.0 * (acf[0] - acf[tau]);
        yin_frame[tau] = diff;
    }

    // Cumulative mean normalized difference
    // d'(tau) = d(tau) / ((1/tau) * sum_{j=1}^{tau} d(j))
    // = d(tau) * tau / cumsum(d)[tau]
    let mut cumsum_d = vec![0.0; n_tau];
    let mut running_d = 0.0;
    for tau in 1..n_tau {
        running_d += yin_frame[tau];
        cumsum_d[tau] = running_d;
    }

    let mut cmnd = vec![0.0; n_tau];
    cmnd[0] = 1.0; // d'(0) = 1 by convention
    for tau in 1..n_tau {
        if cumsum_d[tau] > 0.0 {
            cmnd[tau] = yin_frame[tau] * tau as f64 / cumsum_d[tau];
        } else {
            cmnd[tau] = 1.0;
        }
    }

    // Return only the range [min_period..=max_period]
    cmnd[min_period..=max_period].to_vec()
}

/// Find local minima (troughs) in a slice.
///
/// A trough at k requires strict `<` on the left and non-strict `<=` on the right,
/// matching librosa's convention. This asymmetry means that for a flat valley
/// (e.g. `[0.5, 0.3, 0.3, 0.5]`), the first element of the plateau is the trough.
/// Index 0 is a trough if `x[0] < x[1]` (no left neighbor to compare).
fn find_troughs(x: &[f64]) -> Vec<bool> {
    let n = x.len();
    if n == 0 {
        return vec![];
    }
    let mut troughs = vec![false; n];
    if n >= 2 {
        troughs[0] = x[0] < x[1];
    }
    for k in 1..n.saturating_sub(1) {
        troughs[k] = x[k] < x[k - 1] && x[k] <= x[k + 1];
    }
    troughs
}

/// YIN pitch estimation.
///
/// Matches `librosa.yin(y, fmin=fmin, fmax=fmax, sr=sr, ...)`.
///
/// Returns fundamental frequency (Hz) per frame.
///
/// # Examples
///
/// ```
/// use rosa::pitch::{yin, YinParams};
///
/// let signal = vec![0.0; 8000]; // silence
/// let params = YinParams::new(80.0, 400.0);
/// let f0 = yin(&signal, &params);
/// ```
pub fn yin(y: &[f64], params: &YinParams) -> Vec<f64> {
    let hop_length = params.hop();
    let frame_length = params.frame_length;

    // Pad signal if center
    let padded: Vec<f64>;
    let signal = if params.center {
        let pad = frame_length / 2;
        padded = std::iter::repeat_n(0.0, pad)
            .chain(y.iter().copied())
            .chain(std::iter::repeat_n(0.0, pad))
            .collect();
        &padded
    } else {
        y
    };

    if signal.len() < frame_length {
        return vec![];
    }

    let n_frames = 1 + (signal.len() - frame_length) / hop_length;

    // Period range
    let min_period = (params.sr / params.fmax).floor() as usize;
    let max_period = ((params.sr / params.fmin).ceil() as usize).min(frame_length - 1);

    let mut f0 = Vec::with_capacity(n_frames);

    for i in 0..n_frames {
        let start = i * hop_length;
        let frame = &signal[start..start + frame_length];

        // Compute CMND
        let cmnd = cumulative_mean_normalized_difference(frame, min_period, max_period);
        let shifts = parabolic_interpolation(&cmnd);
        let troughs = find_troughs(&cmnd);

        // Find first trough below threshold
        let mut selected = None;
        for k in 0..cmnd.len() {
            if troughs[k] && cmnd[k] < params.trough_threshold {
                selected = Some(k);
                break;
            }
        }

        // Fallback: global minimum
        let selected = selected.unwrap_or_else(|| {
            cmnd.iter()
                .enumerate()
                .min_by(|(_, a), (_, b)| a.total_cmp(b))
                .map(|(idx, _)| idx)
                .unwrap_or(0)
        });

        // Convert to frequency
        let period = min_period as f64 + selected as f64 + shifts[selected];
        f0.push(params.sr / period);
    }

    f0
}

/// pYIN probabilistic pitch estimation.
///
/// Matches `librosa.pyin(y, fmin=fmin, fmax=fmax, sr=sr, ...)`.
///
/// Returns a [`PyinResult`] with f0, voiced_flag, and voiced_prob.
///
/// # Panics
///
/// Panics if `beta_parameters` contains values that are not both positive.
pub fn pyin(y: &[f64], params: &PyinParams) -> PyinResult {
    use statrs::distribution::{Beta, ContinuousCDF};

    let hop_length = params.hop();
    let frame_length = params.frame_length;

    // Pad signal if center
    let padded: Vec<f64>;
    let signal = if params.center {
        let pad = frame_length / 2;
        padded = std::iter::repeat_n(0.0, pad)
            .chain(y.iter().copied())
            .chain(std::iter::repeat_n(0.0, pad))
            .collect();
        &padded
    } else {
        y
    };

    if signal.len() < frame_length {
        return PyinResult {
            f0: vec![],
            voiced_flag: vec![],
            voiced_prob: vec![],
        };
    }

    let n_frames = 1 + (signal.len() - frame_length) / hop_length;
    let sr = params.sr;

    // Period range
    let min_period = (sr / params.fmax).floor() as usize;
    let max_period = ((sr / params.fmin).ceil() as usize).min(frame_length - 1);
    let n_periods = max_period - min_period + 1;

    // Pitch bin grid
    let n_bins_per_semitone = (1.0 / params.resolution).ceil() as usize;
    let n_pitch_bins = (12.0 * n_bins_per_semitone as f64 * (params.fmax / params.fmin).log2())
        .floor() as usize
        + 1;

    // Beta prior over thresholds.
    // With default (2, 18), this heavily weights low thresholds (high certainty peaks).
    // Uses the `statrs` crate for Beta CDF, matching `scipy.stats.beta.cdf()`.
    // The CDF is evaluated at n_thresholds+1 evenly spaced points in [0, 1] and
    // differenced to get the probability mass in each threshold bin.
    let beta_dist = Beta::new(params.beta_parameters.0, params.beta_parameters.1)
        .expect("invalid beta_parameters: both values must be > 0");
    let thresholds: Vec<f64> = (0..=params.n_thresholds)
        .map(|i| i as f64 / params.n_thresholds as f64)
        .collect();
    let beta_cdf: Vec<f64> = thresholds.iter().map(|&t| beta_dist.cdf(t)).collect();
    let beta_probs: Vec<f64> = beta_cdf.windows(2).map(|w| w[1] - w[0]).collect();

    // Build observation probability matrix [2*n_pitch_bins, n_frames]
    let n_total_states = 2 * n_pitch_bins;
    let mut observation_probs = Matrix::zeros(n_total_states, n_frames);
    let mut voiced_prob_per_frame = vec![0.0; n_frames];

    for frame_idx in 0..n_frames {
        let start = frame_idx * hop_length;
        let frame = &signal[start..start + frame_length];

        // CMND + parabolic interpolation + troughs
        let cmnd = cumulative_mean_normalized_difference(frame, min_period, max_period);
        let shifts = parabolic_interpolation(&cmnd);
        let troughs = find_troughs(&cmnd);

        // Find all trough indices
        let trough_indices: Vec<usize> = (0..cmnd.len()).filter(|&k| troughs[k]).collect();

        if trough_indices.is_empty() {
            // No troughs found — unvoiced frame
            let unvoiced_prob = 1.0 / n_pitch_bins as f64;
            for b in n_pitch_bins..n_total_states {
                *observation_probs.get_mut(b, frame_idx) = unvoiced_prob;
            }
            continue;
        }

        let trough_heights: Vec<f64> = trough_indices.iter().map(|&k| cmnd[k]).collect();

        // For each threshold, find which troughs are below it, apply Boltzmann prior
        let mut yin_probs = vec![0.0; n_periods];

        for (thresh_idx, &threshold) in thresholds[1..].iter().enumerate() {
            // Which troughs are below this threshold?
            let qualifying: Vec<usize> = (0..trough_indices.len())
                .filter(|&ti| trough_heights[ti] < threshold)
                .collect();

            let n_qualifying = qualifying.len();
            if n_qualifying == 0 {
                continue;
            }

            // Boltzmann prior
            let lambda = params.boltzmann_parameter;
            let boltz_unnorm: Vec<f64> = (0..n_qualifying)
                .map(|k| (-(k as f64) / lambda).exp())
                .collect();
            let boltz_sum: f64 = boltz_unnorm.iter().sum();

            for (rank, &ti) in qualifying.iter().enumerate() {
                let boltz_prob = boltz_unnorm[rank] / boltz_sum;
                yin_probs[trough_indices[ti]] += boltz_prob * beta_probs[thresh_idx];
            }
        }

        // Handle no-trough-below-any-threshold fallback
        let global_min_idx = trough_heights
            .iter()
            .enumerate()
            .min_by(|(_, a), (_, b)| a.total_cmp(b))
            .map(|(i, _)| i)
            .unwrap_or(0);

        // Count thresholds below the global min trough height
        let n_below = thresholds[1..]
            .iter()
            .filter(|&&t| t <= trough_heights[global_min_idx])
            .count();
        if n_below > 0 {
            let fallback_prob: f64 = beta_probs[..n_below].iter().sum::<f64>();
            yin_probs[trough_indices[global_min_idx]] += params.no_trough_prob * fallback_prob;
        }

        // Map candidates to pitch bins
        for (period_offset, &prob) in yin_probs.iter().enumerate() {
            if prob <= 0.0 {
                continue;
            }
            let period = min_period as f64 + period_offset as f64 + shifts[period_offset];
            if period <= 0.0 {
                continue;
            }
            let f0_candidate = sr / period;
            if f0_candidate < params.fmin || f0_candidate > params.fmax {
                continue;
            }

            let bin_float = 12.0 * n_bins_per_semitone as f64 * (f0_candidate / params.fmin).log2();
            let bin = bin_float.round().max(0.0).min(n_pitch_bins as f64 - 1.0) as usize;

            *observation_probs.get_mut(bin, frame_idx) += prob;
        }

        // Compute voiced probability for this frame
        let vp: f64 = (0..n_pitch_bins)
            .map(|b| observation_probs.get(b, frame_idx))
            .sum::<f64>()
            .clamp(0.0, 1.0);

        voiced_prob_per_frame[frame_idx] = vp;

        // Fill unvoiced observations
        let unvoiced_val = (1.0 - vp) / n_pitch_bins as f64;
        for b in n_pitch_bins..n_total_states {
            *observation_probs.get_mut(b, frame_idx) = unvoiced_val;
        }
    }

    // Build transition matrix
    let max_semitones_per_frame =
        (params.max_transition_rate * 12.0 * hop_length as f64 / sr).round() as usize;
    let transition_width = max_semitones_per_frame * n_bins_per_semitone + 1;

    let pitch_trans =
        crate::sequence::transition_local(n_pitch_bins, transition_width, "triangle", false);
    let switch_trans = crate::sequence::transition_loop(2, &[1.0 - params.switch_prob]);

    // Kronecker product: switch_trans ⊗ pitch_trans
    let mut full_trans = Matrix::zeros(n_total_states, n_total_states);
    for si in 0..2 {
        for sj in 0..2 {
            let switch_val = switch_trans.get(si, sj);
            for pi in 0..n_pitch_bins {
                for pj in 0..n_pitch_bins {
                    let pitch_val = pitch_trans.get(pi, pj);
                    *full_trans.get_mut(si * n_pitch_bins + pi, sj * n_pitch_bins + pj) =
                        switch_val * pitch_val;
                }
            }
        }
    }

    // Viterbi decoding
    let p_init = vec![1.0 / n_total_states as f64; n_total_states];
    let states = crate::sequence::viterbi(&observation_probs, &full_trans, Some(&p_init));

    // Extract results
    let freqs: Vec<f64> = (0..n_pitch_bins)
        .map(|b| params.fmin * 2.0_f64.powf(b as f64 / (12.0 * n_bins_per_semitone as f64)))
        .collect();

    let mut f0 = Vec::with_capacity(n_frames);
    let mut voiced_flag = Vec::with_capacity(n_frames);

    for &state in &states {
        let pitch_bin = state % n_pitch_bins;
        let is_voiced = state < n_pitch_bins;

        voiced_flag.push(is_voiced);
        if is_voiced {
            f0.push(freqs[pitch_bin]);
        } else {
            f0.push(params.fill_na.unwrap_or(freqs[pitch_bin]));
        }
    }

    PyinResult {
        f0,
        voiced_flag,
        voiced_prob: voiced_prob_per_frame,
    }
}