rosa 0.1.0

//! Viterbi decoding matching `librosa.sequence`.

use crate::matrix::Matrix;

/// Smallest positive normal f64 for log-space safety.
fn tiny() -> f64 {
    f64::MIN_POSITIVE
}

/// Core Viterbi algorithm in log-space.
///
/// - `log_prob`: shape [n_steps, n_states] — log observation likelihoods
/// - `log_trans`: shape [n_states, n_states] — log transition matrix (row-stochastic before log)
/// - `log_p_init`: shape [n_states] — log initial state distribution
///
/// Returns (state_path, log_probability).
fn viterbi_core(
    log_prob: &[f64],
    n_steps: usize,
    n_states: usize,
    log_trans: &[f64],
    log_p_init: &[f64],
) -> (Vec<usize>, f64) {
    let mut value = vec![0.0; n_steps * n_states];
    let mut ptr = vec![0u32; n_steps * n_states];

    // Initialize: value[0, j] = log_prob[0, j] + log_p_init[j]
    for j in 0..n_states {
        value[j] = log_prob[j] + log_p_init[j];
    }

    // Forward pass
    for t in 1..n_steps {
        for j in 0..n_states {
            let mut best_val = f64::NEG_INFINITY;
            let mut best_i = 0u32;

            for i in 0..n_states {
                // trans[i, j] = probability of going from state i to state j
                let v = value[(t - 1) * n_states + i] + log_trans[i * n_states + j];
                if v > best_val {
                    best_val = v;
                    best_i = i as u32;
                }
            }

            ptr[t * n_states + j] = best_i;
            value[t * n_states + j] = log_prob[t * n_states + j] + best_val;
        }
    }

    // Backward pass
    let mut states = vec![0usize; n_steps];
    let last = n_steps - 1;

    // Find best final state
    let mut best_val = f64::NEG_INFINITY;
    for j in 0..n_states {
        let v = value[last * n_states + j];
        if v > best_val {
            best_val = v;
            states[last] = j;
        }
    }

    let logp = best_val;

    // Backtrack
    for t in (0..last).rev() {
        states[t] = ptr[(t + 1) * n_states + states[t + 1]] as usize;
    }

    (states, logp)
}

/// Viterbi decoding for generative (observation likelihood) models.
///
/// Matches `librosa.sequence.viterbi(prob, transition, p_init=p_init)`.
///
/// - `prob`: shape [n_states, n_steps] — `prob[s, t] = P(observation_t | state_s)`
/// - `transition`: shape [n_states, n_states] — row-stochastic transition matrix
/// - `p_init`: optional initial state distribution (default: uniform)
///
/// Returns the most likely state sequence.
///
/// # Panics
///
/// Panics if `transition` is not square with side length `prob.rows()`.
pub fn viterbi(prob: &Matrix, transition: &Matrix, p_init: Option<&[f64]>) -> Vec<usize> {
    let n_states = prob.rows();
    let n_steps = prob.cols();

    assert_eq!(transition.shape(), (n_states, n_states));

    let eps = tiny();

    // Convert to log-space, transposed to [n_steps, n_states]
    let mut log_prob = vec![0.0; n_steps * n_states];
    for t in 0..n_steps {
        for s in 0..n_states {
            log_prob[t * n_states + s] = (prob.get(s, t) + eps).ln();
        }
    }

    let log_trans: Vec<f64> = transition
        .as_slice()
        .iter()
        .map(|&v| (v + eps).ln())
        .collect();

    let default_init = vec![1.0 / n_states as f64; n_states];
    let init = p_init.unwrap_or(&default_init);
    let log_p_init: Vec<f64> = init.iter().map(|&v| (v + eps).ln()).collect();

    let (states, _logp) = viterbi_core(&log_prob, n_steps, n_states, &log_trans, &log_p_init);
    states
}

/// Viterbi decoding for discriminative (posterior) models.
///
/// Matches `librosa.sequence.viterbi_discriminative(prob, transition, p_state, p_init)`.
///
/// - `prob`: shape [n_states, n_steps] — `prob[s, t] = P(state_s | observation_t)`
/// - `transition`: shape [n_states, n_states] — row-stochastic transition matrix
/// - `p_state`: optional marginal state distribution (default: uniform)
/// - `p_init`: optional initial state distribution (default: uniform)
///
/// Applies Bayes' rule to convert to observation likelihoods before decoding.
///
/// # Panics
///
/// Panics if `transition` is not square with side length `prob.rows()`.
pub fn viterbi_discriminative(
    prob: &Matrix,
    transition: &Matrix,
    p_state: Option<&[f64]>,
    p_init: Option<&[f64]>,
) -> Vec<usize> {
    let n_states = prob.rows();
    let n_steps = prob.cols();

    assert_eq!(transition.shape(), (n_states, n_states));

    let eps = tiny();

    let default_state = vec![1.0 / n_states as f64; n_states];
    let state_dist = p_state.unwrap_or(&default_state);

    // Convert discriminative to generative via Bayes' rule:
    // log P(obs|state) ∝ log P(state|obs) - log P(state)
    let log_marginal: Vec<f64> = state_dist.iter().map(|&v| (v + eps).ln()).collect();

    let mut log_prob = vec![0.0; n_steps * n_states];
    for t in 0..n_steps {
        for s in 0..n_states {
            log_prob[t * n_states + s] = (prob.get(s, t) + eps).ln() - log_marginal[s];
        }
    }

    let log_trans: Vec<f64> = transition
        .as_slice()
        .iter()
        .map(|&v| (v + eps).ln())
        .collect();

    let default_init = vec![1.0 / n_states as f64; n_states];
    let init = p_init.unwrap_or(&default_init);
    let log_p_init: Vec<f64> = init.iter().map(|&v| (v + eps).ln()).collect();

    let (states, _logp) = viterbi_core(&log_prob, n_steps, n_states, &log_trans, &log_p_init);
    states
}

/// Viterbi decoding for multi-label binary activation.
///
/// Matches `librosa.sequence.viterbi_binary(prob, transition, p_state, p_init)`.
///
/// - `prob`: shape [n_states, n_steps] — `prob[s, t] = P(state_s active | observation_t)`
/// - `transition`: shape [2, 2] — binary transition matrix applied to all states
/// - `p_state`: marginal probability of each state being active (default: 0.5)
/// - `p_init`: initial probability of each state being active (default: 0.5)
///
/// Each state is decoded independently as a 2-state Viterbi problem.
/// Returns a Matrix of shape [n_states, n_steps] with binary (0/1) values.
///
/// # Panics
///
/// Panics if `transition` is not 2x2.
pub fn viterbi_binary(
    prob: &Matrix,
    transition: &Matrix,
    p_state: Option<f64>,
    p_init: Option<f64>,
) -> Matrix {
    let n_states = prob.rows();
    let n_steps = prob.cols();

    assert_eq!(
        transition.shape(),
        (2, 2),
        "viterbi_binary requires a 2x2 transition matrix"
    );

    let p_s = p_state.unwrap_or(0.5);
    let p_i = p_init.unwrap_or(0.5);

    let state_dist = [1.0 - p_s, p_s];
    let init_dist = [1.0 - p_i, p_i];

    let mut result = Matrix::zeros(n_states, n_steps);

    for s in 0..n_states {
        // Build 2-state probability matrix: [inactive, active]
        let mut prob_binary = Matrix::zeros(2, n_steps);
        for t in 0..n_steps {
            let p = prob.get(s, t);
            *prob_binary.get_mut(0, t) = 1.0 - p;
            *prob_binary.get_mut(1, t) = p;
        }

        let path = viterbi_discriminative(
            &prob_binary,
            transition,
            Some(&state_dist),
            Some(&init_dist),
        );

        for (t, &state) in path.iter().enumerate() {
            *result.get_mut(s, t) = state as f64;
        }
    }

    result
}

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

    #[test]
    fn test_viterbi_2state() {
        // Simple 2-state example from librosa docs
        let p_emit = Matrix::from_vec(vec![0.5, 0.4, 0.1, 0.1, 0.3, 0.6], 2, 3);
        let p_trans = Matrix::from_vec(vec![0.7, 0.3, 0.4, 0.6], 2, 2);
        let p_init = [0.6, 0.4];
        let path = viterbi(&p_emit, &p_trans, Some(&p_init));
        assert_eq!(path, vec![0, 0, 1]);
    }

    #[test]
    fn test_viterbi_uniform_init() {
        // Without explicit p_init (uniform)
        let p_emit = Matrix::from_vec(
            vec![
                0.7, 0.2, 0.1, 0.3, 0.8, 0.2, 0.6, 0.3, 0.5, 0.1, 0.1, 0.2, 0.6, 0.2, 0.1,
            ],
            3,
            5,
        );
        let p_trans = Matrix::from_vec(vec![0.6, 0.3, 0.1, 0.2, 0.5, 0.3, 0.1, 0.2, 0.7], 3, 3);
        let path = viterbi(&p_emit, &p_trans, None);
        // Should decode a reasonable path
        assert_eq!(path.len(), 5);
    }
}