rosa 0.1.0

use crate::matrix::Matrix;

/// Parameters for computing delta (derivative) features.
pub struct DeltaParams {
    /// Window length for the Savitzky-Golay filter. Must be odd and >= 3.
    pub width: usize,
    /// Derivative order. 1 = first derivative, 2 = second derivative.
    pub order: usize,
}

impl Default for DeltaParams {
    fn default() -> Self {
        Self { width: 9, order: 1 }
    }
}

/// Compute delta (Savitzky-Golay derivative) features along the time axis.
///
/// Matches `librosa.feature.delta` which calls
/// `scipy.signal.savgol_filter(data, width, deriv=order, mode='interp')`.
///
/// Input:  `data [n_features, T]`
/// Output: `[n_features, T]` — derivative along the column (time) axis
pub fn delta(data: &Matrix, params: &DeltaParams) -> Matrix {
    let width = params.width;
    let order = params.order;
    assert!(width >= 3 && width % 2 == 1, "width must be odd and >= 3");
    assert!(order >= 1, "order must be >= 1");

    let n_rows = data.rows();
    let t = data.cols();
    let half = (width - 1) / 2;

    // Precompute all coefficient vectors for each column position.
    // Positions 0..half and (t-half)..t use asymmetric windows (mode='interp').
    // Interior positions all use the same centered-window coefficients.

    // For position i (0-indexed), the window into the data is:
    //   left boundary  (i < half):      data cols [0, width)
    //   interior       (half <= i < t-half): data cols [i-half, i+half+1)
    //   right boundary (i >= t-half):   data cols [t-width, t)
    //
    // Within that window, eval_pos is the index of column i in the window:
    //   left boundary:   eval_pos = i
    //   interior:        eval_pos = half
    //   right boundary:  eval_pos = i - (t - width)

    // Compute interior coefficients once
    let interior_coeffs = sg_coeffs(width, order, half as f64);

    // Precompute boundary coefficient sets
    let left_coeffs: Vec<Vec<f64>> = (0..half)
        .map(|i| sg_coeffs(width, order, i as f64))
        .collect();
    let right_coeffs: Vec<Vec<f64>> = (0..half)
        .map(|k| {
            // column index = t - half + k, eval_pos in window = (t - half + k) - (t - width)
            //              = width - half + k - 1 = half + 1 + k
            let eval_pos = (half + 1 + k) as f64;
            sg_coeffs(width, order, eval_pos)
        })
        .collect();

    let mut result = Matrix::zeros(n_rows, t);

    for r in 0..n_rows {
        for c in 0..t {
            let val = if c < half {
                // Left boundary: window [0, width)
                let coeffs = &left_coeffs[c];
                (0..width).map(|j| coeffs[j] * data.get(r, j)).sum()
            } else if c >= t.saturating_sub(half) {
                // Right boundary: window [t-width, t)
                let k = c - (t - half);
                let coeffs = &right_coeffs[k];
                let win_start = t - width;
                (0..width)
                    .map(|j| coeffs[j] * data.get(r, win_start + j))
                    .sum()
            } else {
                // Interior: window [c-half, c+half+1)
                let win_start = c - half;
                (0..width)
                    .map(|j| interior_coeffs[j] * data.get(r, win_start + j))
                    .sum()
            };
            *result.get_mut(r, c) = val;
        }
    }

    result
}

/// Compute Savitzky-Golay filter coefficients for a window of `window_length` samples,
/// evaluating the `deriv`-th derivative at position `eval_pos` within the window.
///
/// Uses Vandermonde matrix + normal equations (least-squares).
/// `polyorder` equals `deriv` (minimum polynomial order for non-trivial derivative).
fn sg_coeffs(window_length: usize, deriv: usize, eval_pos: f64) -> Vec<f64> {
    // polyorder = deriv ensures we can represent the derivative we need.
    // But we need polyorder >= deriv; use max(deriv, 1) as minimum to avoid trivial filter.
    // librosa uses polyorder = order (same as deriv) by default.
    let polyorder = deriv;
    let p = polyorder + 1; // number of polynomial coefficients
    let n = window_length;

    // Build Vandermonde matrix A[j,k] = (j - eval_pos)^k
    let mut a = vec![0.0f64; n * p];
    for j in 0..n {
        let x = j as f64 - eval_pos;
        let mut xpow = 1.0f64;
        for k in 0..p {
            a[j * p + k] = xpow;
            xpow *= x;
        }
    }

    // Compute A^T A  (p x p)
    let mut ata = vec![0.0f64; p * p];
    for k1 in 0..p {
        for k2 in 0..p {
            let sum: f64 = (0..n).map(|j| a[j * p + k1] * a[j * p + k2]).sum();
            ata[k1 * p + k2] = sum;
        }
    }

    // Solve (A^T A) v = e_deriv  via Gaussian elimination with partial pivoting
    let mut v = vec![0.0f64; p];
    v[deriv] = 1.0;
    gaussian_solve(&ata, &mut v, p);

    // Filter weights: w[j] = factorial(deriv) * sum_k v[k] * A[j,k]
    let fact = factorial(deriv) as f64;
    (0..n)
        .map(|j| fact * (0..p).map(|k| v[k] * a[j * p + k]).sum::<f64>())
        .collect()
}

/// Solve `A x = b` in-place (result in `b`) using Gaussian elimination with partial pivoting.
/// `a` is row-major `n x n`, `b` is length `n`.
fn gaussian_solve(a: &[f64], b: &mut [f64], n: usize) {
    let mut mat = a.to_vec(); // working copy
    let mut rhs = b.to_vec();

    for col in 0..n {
        // Find pivot
        let pivot_row = (col..n).max_by(|&r1, &r2| {
            mat[r1 * n + col]
                .abs()
                .partial_cmp(&mat[r2 * n + col].abs())
                .unwrap()
        });
        let pivot_row = pivot_row.unwrap();

        // Swap rows
        if pivot_row != col {
            for j in 0..n {
                mat.swap(col * n + j, pivot_row * n + j);
            }
            rhs.swap(col, pivot_row);
        }

        let pivot = mat[col * n + col];
        if pivot.abs() < 1e-14 {
            continue; // singular or near-singular
        }

        // Eliminate below
        for row in (col + 1)..n {
            let factor = mat[row * n + col] / pivot;
            for j in col..n {
                let v = mat[col * n + j];
                mat[row * n + j] -= factor * v;
            }
            rhs[row] -= factor * rhs[col];
        }
    }

    // Back substitution
    for i in (0..n).rev() {
        let mut sum = rhs[i];
        for j in (i + 1)..n {
            sum -= mat[i * n + j] * b[j];
        }
        b[i] = sum / mat[i * n + i];
    }
}

fn factorial(n: usize) -> usize {
    (1..=n).product()
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_delta_shape() {
        let data = Matrix::from_fn(20, 100, |r, c| (r + c) as f64);
        let result = delta(&data, &DeltaParams::default());
        assert_eq!(result.shape(), (20, 100));
    }

    #[test]
    fn test_delta_linear_signal() {
        // For a linear signal, first derivative should be constant (= slope).
        // delta with width=9, order=1 on y = c*x should give ~c everywhere.
        let slope = 2.0;
        let data = Matrix::from_fn(1, 50, |_r, c| slope * c as f64);
        let result = delta(&data, &DeltaParams::default());
        for c in 0..50 {
            let diff = (result.get(0, c) - slope).abs();
            assert!(
                diff < 1e-9,
                "col {c}: got {}, expected {slope}",
                result.get(0, c)
            );
        }
    }

    #[test]
    fn test_delta_interior_coefficients_sum_to_zero() {
        // SG derivative filter coefficients for interior should sum to 0 (no DC)
        let coeffs = sg_coeffs(9, 1, 4.0);
        let sum: f64 = coeffs.iter().sum();
        assert!(sum.abs() < 1e-12, "sum = {sum}");
    }
}