rosa 0.1.0

//! Transition matrix builders matching `librosa.sequence`.

use crate::matrix::Matrix;

/// Uniform transition matrix.
///
/// Matches `librosa.sequence.transition_uniform(n_states)`.
/// Every cell is `1.0 / n_states`.
pub fn transition_uniform(n_states: usize) -> Matrix {
    let val = 1.0 / n_states as f64;
    Matrix::from_vec(vec![val; n_states * n_states], n_states, n_states)
}

/// Self-loop transition matrix.
///
/// Matches `librosa.sequence.transition_loop(n_states, prob)`.
///
/// Diagonal = `prob[i]`, off-diagonal = `(1 - prob[i]) / (n_states - 1)`.
/// If `prob` has one element, it is broadcast to all states.
///
/// # Panics
///
/// Panics if `n_states <= 1` or if `prob.len()` is not 1 or `n_states`.
pub fn transition_loop(n_states: usize, prob: &[f64]) -> Matrix {
    assert!(n_states > 1, "transition_loop requires n_states > 1");

    let prob_vec: Vec<f64> = if prob.len() == 1 {
        vec![prob[0]; n_states]
    } else {
        assert_eq!(
            prob.len(),
            n_states,
            "prob must have length 1 or n_states ({})",
            n_states
        );
        prob.to_vec()
    };

    Matrix::from_fn(n_states, n_states, |i, j| {
        if i == j {
            prob_vec[i]
        } else {
            (1.0 - prob_vec[i]) / (n_states - 1) as f64
        }
    })
}

/// Localized transition matrix with windowed transitions.
///
/// Matches `librosa.sequence.transition_local(n_states, width, window, wrap)`.
///
/// Each state can only transition to nearby states within `width` bins.
/// The transition probabilities follow the specified window shape.
///
/// Supported window types: `"triangle"`, `"ones"`, `"rectangular"`, `"boxcar"`,
/// `"hann"`, `"hamming"`.
///
/// # Panics
///
/// Panics if `n_states <= 1` or if `window` is not a supported type.
pub fn transition_local(n_states: usize, width: usize, window: &str, wrap: bool) -> Matrix {
    assert!(n_states > 1, "transition_local requires n_states > 1");

    let width = width.max(1);

    // Generate window of the given type
    let win = match window {
        "triangle" => crate::windows::triangular(width, false),
        "ones" | "rectangular" | "boxcar" => vec![1.0; width],
        "hann" => crate::windows::hann(width, false),
        "hamming" => crate::windows::hamming(width, false),
        _ => panic!("Unsupported window type: {window}"),
    };

    // Pad window to n_states length (center-padded)
    let padded = crate::dsp::util::pad_center(&win, n_states);

    let mut data = vec![0.0; n_states * n_states];

    for i in 0..n_states {
        // Roll the padded window so it's centered on state i
        // We need to shift by (n_states/2 + i + 1) to center on state i
        // This matches librosa: np.roll(trans_row, n_states // 2 + i + 1)
        let shift = (n_states / 2 + i + 1) % n_states;

        for j in 0..n_states {
            let src = (j + n_states - shift) % n_states;
            data[i * n_states + j] = padded[src];
        }

        if !wrap {
            // Zero out elements beyond local region
            let half = width / 2;
            let lo = i.saturating_sub(half);
            let hi = (i + half + width % 2).min(n_states);

            for j in 0..n_states {
                if j < lo || j >= hi {
                    data[i * n_states + j] = 0.0;
                }
            }
        }

        // Normalize row to sum to 1
        let row_start = i * n_states;
        let row_sum: f64 = data[row_start..row_start + n_states].iter().sum();
        if row_sum > 0.0 {
            for j in 0..n_states {
                data[row_start + j] /= row_sum;
            }
        }
    }

    Matrix::from_vec(data, n_states, n_states)
}

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

    #[allow(dead_code)]
    fn assert_matrix_approx(actual: &Matrix, expected: &[f64], tol: f64, name: &str) {
        let actual_data = actual.as_slice();
        assert_eq!(
            actual_data.len(),
            expected.len(),
            "{name}: size mismatch: {} vs {}",
            actual_data.len(),
            expected.len()
        );
        for (i, (a, e)) in actual_data.iter().zip(expected.iter()).enumerate() {
            let diff = (a - e).abs();
            assert!(
                diff < tol,
                "{name}[{i}]: got {a}, expected {e}, diff {diff}"
            );
        }
    }

    #[test]
    fn test_transition_uniform() {
        let t = transition_uniform(3);
        assert_eq!(t.shape(), (3, 3));
        let v = 1.0 / 3.0;
        for i in 0..3 {
            for j in 0..3 {
                assert!((t.get(i, j) - v).abs() < 1e-15);
            }
        }
    }

    #[test]
    fn test_transition_loop_scalar() {
        let t = transition_loop(3, &[0.5]);
        assert_eq!(t.shape(), (3, 3));
        for i in 0..3 {
            assert!((t.get(i, i) - 0.5).abs() < 1e-15);
            for j in 0..3 {
                if i != j {
                    assert!((t.get(i, j) - 0.25).abs() < 1e-15);
                }
            }
        }
    }

    #[test]
    fn test_transition_loop_per_state() {
        let t = transition_loop(3, &[0.8, 0.5, 0.25]);
        assert!((t.get(0, 0) - 0.8).abs() < 1e-15);
        assert!((t.get(0, 1) - 0.1).abs() < 1e-15);
        assert!((t.get(1, 1) - 0.5).abs() < 1e-15);
        assert!((t.get(2, 2) - 0.25).abs() < 1e-15);
        assert!((t.get(2, 0) - 0.375).abs() < 1e-15);
    }

    #[test]
    fn test_transition_local_rows_sum_to_one() {
        let t = transition_local(5, 3, "triangle", false);
        for i in 0..5 {
            let row_sum: f64 = (0..5).map(|j| t.get(i, j)).sum();
            assert!(
                (row_sum - 1.0).abs() < 1e-12,
                "row {i} sums to {row_sum}, not 1.0"
            );
        }
    }

    #[test]
    fn test_transition_local_wrap_rows_sum_to_one() {
        let t = transition_local(5, 3, "triangle", true);
        for i in 0..5 {
            let row_sum: f64 = (0..5).map(|j| t.get(i, j)).sum();
            assert!(
                (row_sum - 1.0).abs() < 1e-12,
                "row {i} sums to {row_sum}, not 1.0"
            );
        }
    }
}