rosa 0.1.0

use crate::matrix::Matrix;
use crate::windows::hann;

/// Parameters for `path_enhance`.
pub struct PathEnhanceParams {
    /// Filter length (number of samples in each diagonal kernel).
    pub n: usize,
    /// Maximum slope ratio. Default 2.0.
    pub max_ratio: f64,
    /// Minimum slope ratio. `None` → `1 / max_ratio`.
    pub min_ratio: Option<f64>,
    /// Number of diagonal filters. Default 7.
    pub n_filters: usize,
    /// If true, subtract the mean from each kernel. Default false.
    pub zero_mean: bool,
    /// If true, clamp output to >= 0. Default true.
    pub clip: bool,
}

impl Default for PathEnhanceParams {
    fn default() -> Self {
        Self {
            n: 51,
            max_ratio: 2.0,
            min_ratio: None,
            n_filters: 7,
            zero_mean: false,
            clip: true,
        }
    }
}

/// Enhance diagonal paths in a recurrence/similarity matrix.
///
/// Matches `librosa.segment.path_enhance`.
///
/// Input:  `rec [T, T]`
/// Output: `[T, T]` — enhanced matrix
pub fn path_enhance(rec: &Matrix, params: &PathEnhanceParams) -> Matrix {
    let t = rec.rows();
    assert_eq!(t, rec.cols(), "recurrence matrix must be square");

    let min_ratio = params.min_ratio.unwrap_or(1.0 / params.max_ratio);
    let max_ratio = params.max_ratio;
    let n_filters = params.n_filters;

    // Generate n_filters ratios log-spaced between min_ratio and max_ratio
    let ratios: Vec<f64> = if n_filters == 1 {
        vec![(min_ratio * max_ratio).sqrt()]
    } else {
        let log_min = min_ratio.ln();
        let log_max = max_ratio.ln();
        (0..n_filters)
            .map(|i| {
                let t = i as f64 / (n_filters - 1) as f64;
                (log_min + t * (log_max - log_min)).exp()
            })
            .collect()
    };

    let window = hann(params.n, false); // symmetric window

    let mut output = vec![f64::NEG_INFINITY; t * t];

    for &ratio in &ratios {
        let kernel = diagonal_filter(&window, ratio, params.zero_mean);
        let convolved = convolve2d_reflect(rec, &kernel);

        for idx in 0..t * t {
            if convolved[idx] > output[idx] {
                output[idx] = convolved[idx];
            }
        }
    }

    // Replace NEG_INFINITY with 0 (no filter touched this cell — shouldn't happen)
    for v in output.iter_mut() {
        if v.is_infinite() && *v < 0.0 {
            *v = 0.0;
        }
    }

    if params.clip {
        for v in output.iter_mut() {
            if *v < 0.0 {
                *v = 0.0;
            }
        }
    }

    Matrix::from_vec(output, t, t)
}

/// Build a 2D diagonal filter kernel from a 1D window with the given slope.
///
/// The kernel has `n` non-zero samples placed along a line with slope `ratio`.
/// Returns a flat row-major 2D array of shape `[height, width]`.
fn diagonal_filter(window: &[f64], ratio: f64, zero_mean: bool) -> (Vec<f64>, usize, usize) {
    let n = window.len();

    // Sample n points along a diagonal with slope `ratio`.
    // The diagonal goes from (0,0) to ((n-1)*ratio, n-1) in (row, col) space,
    // i.e., for step k: row_f = k * ratio, col_f = k * 1.0
    // Round to integer offsets and build a sparse 2D kernel.

    let mut row_offsets: Vec<isize> = (0..n)
        .map(|k| (k as f64 * ratio).round() as isize)
        .collect();
    let col_offsets: Vec<isize> = (0..n).map(|k| k as isize).collect();

    // Shift so minimum row offset = 0
    let min_row = *row_offsets.iter().min().unwrap();
    for r in row_offsets.iter_mut() {
        *r -= min_row;
    }
    let max_row = *row_offsets.iter().max().unwrap() as usize;
    let max_col = *col_offsets.iter().max().unwrap() as usize;

    let height = max_row + 1;
    let width = max_col + 1;

    let mut kernel = vec![0.0f64; height * width];
    for k in 0..n {
        let r = row_offsets[k] as usize;
        let c = col_offsets[k] as usize;
        kernel[r * width + c] = window[k];
    }

    if zero_mean {
        let mean: f64 = kernel.iter().sum::<f64>() / (height * width) as f64;
        for v in kernel.iter_mut() {
            *v -= mean;
        }
    }

    (kernel, height, width)
}

/// 2D convolution with reflect padding (matching scipy.ndimage.convolve semantics).
///
/// `kernel` is (height x width) row-major, returns flat [t x t] result.
fn convolve2d_reflect(rec: &Matrix, kernel_tuple: &(Vec<f64>, usize, usize)) -> Vec<f64> {
    let (kernel, kh, kw) = kernel_tuple;
    let t = rec.rows();
    let ph = kh / 2; // padding rows
    let pw = kw / 2; // padding cols

    let mut result = vec![0.0f64; t * t];

    // For each output (r, c), compute sum over kernel
    for r in 0..t {
        for c in 0..t {
            let mut sum = 0.0f64;
            for kr in 0..*kh {
                for kc in 0..*kw {
                    // Source position (before padding): (r + kr - ph, c + kc - pw)
                    let sr = reflect_index(r as isize + kr as isize - ph as isize, t);
                    let sc = reflect_index(c as isize + kc as isize - pw as isize, t);
                    sum += rec.get(sr, sc) * kernel[kr * kw + kc];
                }
            }
            result[r * t + c] = sum;
        }
    }

    result
}

/// Reflect boundary index clamping for a signal of length `n`.
/// Matches numpy/scipy `mode='reflect'` (edge values reflected inward).
fn reflect_index(i: isize, n: usize) -> usize {
    if n == 0 {
        return 0;
    }
    let n = n as isize;
    let mut i = i;
    // Fold into [0, n)
    if i < 0 {
        i = -i - 1;
    }
    if i >= n {
        i = 2 * n - i - 1;
    }
    // Clamp for safety
    i.clamp(0, n - 1) as usize
}