arcweight 0.3.0

//! $`O(V+E)`$ shortest distance for acyclic FSTs.
//!
//! This module provides specialized shortest distance algorithms for acyclic (DAG)
//! FSTs, achieving optimal $`O(V+E)`$ complexity by exploiting the topological structure.
//! Since most FSTs in speech/NLP applications are acyclic, this is the preferred
//! algorithm when applicable.
//!
//! # Complexity Comparison
//!
//! | Algorithm | Time | Space | Requirements |
//! |-----------|------|-------|--------------|
//! | Acyclic (this) | $`O(V+E)`$ | $`O(V)`$ | Acyclic FST |
//! | Dijkstra | $`O((V+E)\log V)`$ | $`O(V)`$ | Non-negative weights |
//! | Bellman-Ford | $`O(VE)`$ | $`O(V)`$ | General |
//!
//! # Algorithm
//!
//! The acyclic algorithm uses Kahn's topological sort followed by single-pass relaxation:
//!
//! 1. **Topological Sort:** Order states so all arcs go forward — $`O(V+E)`$
//! 2. **Relaxation:** Process states in order, relax outgoing arcs — $`O(V+E)`$
//!
//! For state $`q`$ in topological order with distance $`d[q]`$:
//! ```text
//! for each arc q --w--> r:
//!     d[r] ← d[r] ⊕ (d[q] ⊗ w)
//! ```
//!
//! # When to Use
//!
//! Most FSTs in practice are acyclic:
//!
//! - **Lexicon FSTs:** Dictionary lookups
//! - **Pronunciation models:** Grapheme-to-phoneme mappings
//! - **N-gram LMs:** Finite-order language models
//! - **Linear-chain CRFs:** Sequence labeling models
//! - **Morphological analyzers:** Word decomposition
//!
//! Use [`is_acyclic`] to check, or [`shortest_distance_auto`] for automatic selection.
//!
//! # Example
//!
//! ```rust
//! use arcweight::prelude::*;
//! use arcweight::algorithms::{shortest_distance_acyclic, is_acyclic};
//!
//! let mut fst = VectorFst::<TropicalWeight>::new();
//! let s0 = fst.add_state();
//! let s1 = fst.add_state();
//! let s2 = fst.add_state();
//! fst.set_start(s0);
//! fst.set_final(s2, TropicalWeight::one());
//!
//! // Two paths to s2
//! fst.add_arc(s0, Arc::new(1, 1, TropicalWeight::new(3.0), s2));  // Direct: 3.0
//! fst.add_arc(s0, Arc::new(2, 2, TropicalWeight::new(1.0), s1));  // Via s1: 1.0 + 0.5 = 1.5
//! fst.add_arc(s1, Arc::new(3, 3, TropicalWeight::new(0.5), s2));
//!
//! // Verify acyclic and compute distances
//! assert!(is_acyclic(&fst));
//! let d = shortest_distance_acyclic(&fst)?;
//!
//! assert_eq!(*d[0].value(), 0.0);  // Start
//! assert_eq!(*d[1].value(), 1.0);  // Via s0->s1
//! assert_eq!(*d[2].value(), 1.5);  // min(3.0, 1.0+0.5) = 1.5
//! # Ok::<(), arcweight::Error>(())
//! ```
//!
//! # References
//!
//! \[1\] Mohri, M. 2002. Semiring frameworks and algorithms for shortest-distance
//!     problems. *Journal of Automata, Languages and Combinatorics* 7, 3, 321-350.
//!
//! \[2\] Cormen, T. H., Leiserson, C. E., Rivest, R. L., and Stein, C. 2009.
//!     *Introduction to Algorithms* (3rd ed.). MIT Press. Chapter 24: Single-Source
//!     Shortest Paths.
//!
//! \[3\] Kahn, A. B. 1962. Topological sorting of large networks. *Communications
//!     of the ACM* 5, 11 (November 1962), 558-562. DOI: <https://doi.org/10.1145/368996.369025>

use crate::fst::{Fst, StateId};
use crate::semiring::Semiring;
use crate::{Error, Result};

/// Computes shortest distances from the start state for acyclic FSTs.
///
/// Uses topological sorting to process states in dependency order, achieving
/// optimal $`O(V+E)`$ time complexity. This is the preferred algorithm for
/// acyclic FSTs, which are common in speech and NLP applications.
///
/// # Type Parameters
///
/// * `W` - Weight type implementing [`Semiring`]
/// * `F` - FST type implementing [`Fst`]
///
/// # Arguments
///
/// * `fst` - Input FST (must be acyclic)
///
/// # Returns
///
/// Vector of weights where `result[s]` is the shortest distance from the start
/// state to state `s`. Returns `W::zero()` (semiring zero) for unreachable states.
///
/// # Complexity
///
/// - **Time:** $`O(V + E)`$
/// - **Space:** $`O(V)`$
///
/// # Algorithm
///
/// 1. Compute topological order using Kahn's algorithm
/// 2. Initialize $`d[q_0] = \bar{1}`$, $`d[q] = \bar{0}`$ otherwise
/// 3. For each state $`q`$ in topological order:
///    - For each arc $`(q, w, r)`$: $`d[r] \leftarrow d[r] \oplus (d[q] \otimes w)`$
///
/// # Errors
///
/// Returns [`Error::InvalidOperation`] if the FST
/// contains cycles (detected during topological sort).
///
/// # Example
///
/// ```rust
/// use arcweight::prelude::*;
/// use arcweight::algorithms::shortest_distance_acyclic;
///
/// let mut fst = VectorFst::<TropicalWeight>::new();
/// let s0 = fst.add_state();
/// let s1 = fst.add_state();
/// fst.set_start(s0);
/// fst.set_final(s1, TropicalWeight::one());
/// fst.add_arc(s0, Arc::new(1, 1, TropicalWeight::new(2.5), s1));
///
/// let distances = shortest_distance_acyclic(&fst)?;
/// assert_eq!(*distances[0].value(), 0.0);  // Start state (tropical one)
/// assert_eq!(*distances[1].value(), 2.5);  // Via single arc
/// # Ok::<(), arcweight::Error>(())
/// ```
///
/// # References
///
/// \[1\] Mohri, M. 2002. Semiring frameworks and algorithms for shortest-distance
///     problems. *Journal of Automata, Languages and Combinatorics* 7, 3, 321-350.
///
/// # See Also
///
/// - [`shortest_distance`](crate::algorithms::shortest_distance()) - General algorithm for cyclic FSTs
/// - [`is_acyclic`] - Check if FST is acyclic
/// - [`shortest_distance_auto`] - Automatic algorithm selection
pub fn shortest_distance_acyclic<W, F>(fst: &F) -> Result<Vec<W>>
where
    W: Semiring,
    F: Fst<W>,
{
    let num_states = fst.num_states();
    if num_states == 0 {
        return Ok(Vec::new());
    }

    let start = match fst.start() {
        Some(s) => s,
        None => return Ok(vec![W::zero(); num_states]),
    };

    // Compute topological order using Kahn's algorithm
    let topo_order = topological_sort(fst)?;

    // Initialize distances
    let mut distances = vec![W::zero(); num_states];
    distances[start as usize] = W::one();

    // Process states in topological order
    for &state in &topo_order {
        let dist = distances[state as usize].clone();
        // Skip unreachable states: if distance is the semiring zero (additive identity),
        // no paths can reach this state. For tropical/log semirings, zero = infinity.
        // This optimization is safe because zero.times(w) = zero for any weight w.
        if Semiring::is_zero(&dist) {
            continue;
        }

        // Relax all outgoing arcs
        for arc in fst.arcs(state) {
            let new_dist = dist.times(&arc.weight);
            distances[arc.nextstate as usize] = distances[arc.nextstate as usize].plus(&new_dist);
        }
    }

    Ok(distances)
}

/// Computes reverse shortest distances (to final states) for acyclic FSTs.
///
/// Computes the shortest distance from each state to any final state in
/// $`O(V+E)`$ time. This is the "backward" version of shortest distance,
/// useful for:
///
/// - Weight pushing (toward final states)
/// - A* search heuristics
/// - Pruning based on total path weights
///
/// # Type Parameters
///
/// * `W` - Weight type implementing [`Semiring`]
/// * `F` - FST type implementing [`Fst`]
///
/// # Arguments
///
/// * `fst` - Input FST (must be acyclic)
///
/// # Returns
///
/// Vector of weights where `result[s]` is the shortest distance from state
/// `s` to any final state (including final weights).
///
/// # Complexity
///
/// - **Time:** $`O(V + E)`$
/// - **Space:** $`O(V + E)`$ (for reverse adjacency list)
///
/// # Algorithm
///
/// 1. Build reverse adjacency list
/// 2. Compute topological order and reverse it
/// 3. Initialize with final weights
/// 4. Process states in reverse topological order, propagating distances backward
///
/// # Example
///
/// ```rust
/// use arcweight::prelude::*;
/// use arcweight::algorithms::shortest_distance_acyclic_reverse;
///
/// let mut fst = VectorFst::<TropicalWeight>::new();
/// let s0 = fst.add_state();
/// let s1 = fst.add_state();
/// let s2 = fst.add_state();
/// fst.set_start(s0);
/// fst.set_final(s2, TropicalWeight::new(0.5));
///
/// fst.add_arc(s0, Arc::new(1, 1, TropicalWeight::new(1.0), s1));
/// fst.add_arc(s1, Arc::new(2, 2, TropicalWeight::new(2.0), s2));
///
/// let d = shortest_distance_acyclic_reverse(&fst)?;
/// // d[s2] = 0.5 (final weight)
/// // d[s1] = 2.0 + 0.5 = 2.5
/// // d[s0] = 1.0 + 2.0 + 0.5 = 3.5
/// # Ok::<(), arcweight::Error>(())
/// ```
pub fn shortest_distance_acyclic_reverse<W, F>(fst: &F) -> Result<Vec<W>>
where
    W: Semiring,
    F: Fst<W>,
{
    let num_states = fst.num_states();
    if num_states == 0 {
        return Ok(Vec::new());
    }

    // Compute topological order and reverse it
    let mut topo_order = topological_sort(fst)?;
    topo_order.reverse();

    // Initialize distances from final states
    let mut distances = vec![W::zero(); num_states];
    for state in fst.states() {
        if let Some(w) = fst.final_weight(state) {
            distances[state as usize] = w.clone();
        }
    }

    // Build reverse adjacency list for efficient backward traversal
    let mut incoming: Vec<Vec<(StateId, W)>> = vec![Vec::new(); num_states];
    for state in fst.states() {
        for arc in fst.arcs(state) {
            incoming[arc.nextstate as usize].push((state, arc.weight.clone()));
        }
    }

    // Process states in reverse topological order
    for &state in &topo_order {
        let dist = distances[state as usize].clone();
        if Semiring::is_zero(&dist) {
            continue;
        }

        // Propagate distances backward
        for (prev_state, weight) in &incoming[state as usize] {
            let new_dist = weight.times(&dist);
            distances[*prev_state as usize] = distances[*prev_state as usize].plus(&new_dist);
        }
    }

    Ok(distances)
}

/// Computes topological order using Kahn's algorithm.
///
/// Returns states in topological order where all arcs go from earlier to
/// later states. Returns an error if the FST contains cycles.
///
/// # Complexity
///
/// - **Time:** $`O(V + E)`$
/// - **Space:** $`O(V)`$
///
/// # References
///
/// \[1\] Kahn, A. B. 1962. Topological sorting of large networks. *Communications
///     of the ACM* 5, 11 (November 1962), 558-562. DOI: <https://doi.org/10.1145/368996.369025>
fn topological_sort<W, F>(fst: &F) -> Result<Vec<StateId>>
where
    W: Semiring,
    F: Fst<W>,
{
    let num_states = fst.num_states();
    if num_states == 0 {
        return Ok(Vec::new());
    }

    // Compute in-degrees
    let mut in_degree = vec![0usize; num_states];
    for state in fst.states() {
        for arc in fst.arcs(state) {
            in_degree[arc.nextstate as usize] += 1;
        }
    }

    // Initialize queue with states that have no incoming edges
    let mut queue: Vec<StateId> = Vec::new();
    for state in fst.states() {
        if in_degree[state as usize] == 0 {
            queue.push(state);
        }
    }

    // Process queue
    let mut result = Vec::with_capacity(num_states);
    while let Some(state) = queue.pop() {
        result.push(state);

        for arc in fst.arcs(state) {
            in_degree[arc.nextstate as usize] -= 1;
            if in_degree[arc.nextstate as usize] == 0 {
                queue.push(arc.nextstate);
            }
        }
    }

    // Check for cycles
    if result.len() != num_states {
        return Err(Error::InvalidOperation(
            "FST contains cycles - use general shortest_distance instead".to_string(),
        ));
    }

    Ok(result)
}

/// Checks if an FST is acyclic (contains no cycles).
///
/// Uses topological sort to detect cycles in $`O(V+E)`$ time. This can be used
/// to determine whether [`shortest_distance_acyclic`] is applicable.
///
/// # Example
///
/// ```rust
/// use arcweight::prelude::*;
/// use arcweight::algorithms::is_acyclic;
///
/// // Acyclic FST
/// let mut fst = VectorFst::<TropicalWeight>::new();
/// let s0 = fst.add_state();
/// let s1 = fst.add_state();
/// fst.set_start(s0);
/// fst.add_arc(s0, Arc::new(1, 1, TropicalWeight::one(), s1));
/// assert!(is_acyclic(&fst));
///
/// // Cyclic FST
/// fst.add_arc(s1, Arc::new(2, 2, TropicalWeight::one(), s0)); // Creates cycle
/// assert!(!is_acyclic(&fst));
/// ```
pub fn is_acyclic<W, F>(fst: &F) -> bool
where
    W: Semiring,
    F: Fst<W>,
{
    topological_sort(fst).is_ok()
}

/// Automatically selects the optimal shortest distance algorithm.
///
/// Checks if the FST is acyclic and uses the $`O(V+E)`$ algorithm if possible,
/// otherwise falls back to the general iterative algorithm for cyclic FSTs.
///
/// This is the recommended entry point when the FST structure is unknown.
///
/// # Example
///
/// ```rust
/// use arcweight::prelude::*;
/// use arcweight::algorithms::shortest_distance_auto;
///
/// let mut fst = VectorFst::<TropicalWeight>::new();
/// let s0 = fst.add_state();
/// let s1 = fst.add_state();
/// fst.set_start(s0);
/// fst.set_final(s1, TropicalWeight::one());
/// fst.add_arc(s0, Arc::new(1, 1, TropicalWeight::new(1.0), s1));
///
/// // Automatically uses acyclic algorithm for this FST
/// let distances = shortest_distance_auto(&fst)?;
/// # Ok::<(), arcweight::Error>(())
/// ```
///
/// # See Also
///
/// - [`shortest_distance_acyclic`] - Explicit acyclic algorithm
/// - [`shortest_distance`](crate::algorithms::shortest_distance()) - General algorithm
/// - [`is_acyclic`] - Check FST structure
pub fn shortest_distance_auto<W, F>(fst: &F) -> Result<Vec<W>>
where
    W: Semiring,
    F: Fst<W>,
{
    // Try acyclic algorithm first (will fail if cycles exist)
    match shortest_distance_acyclic(fst) {
        Ok(distances) => Ok(distances),
        Err(_) => {
            // Fall back to general algorithm
            crate::algorithms::shortest_distance(fst)
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::fst::{MutableFst, VectorFst};
    use crate::semiring::{Semiring, TropicalWeight};
    use num_traits::One;

    #[test]
    fn test_acyclic_shortest_distance_simple() {
        let mut fst = VectorFst::<TropicalWeight>::new();
        let s0 = fst.add_state();
        let s1 = fst.add_state();
        let s2 = fst.add_state();

        fst.set_start(s0);
        fst.set_final(s2, TropicalWeight::one());

        // Two paths: s0 -> s2 (weight 3.0) or s0 -> s1 -> s2 (weight 1.0 + 0.5 = 1.5)
        fst.add_arc(s0, crate::arc::Arc::new(1, 1, TropicalWeight::new(3.0), s2));
        fst.add_arc(s0, crate::arc::Arc::new(2, 2, TropicalWeight::new(1.0), s1));
        fst.add_arc(s1, crate::arc::Arc::new(3, 3, TropicalWeight::new(0.5), s2));

        let distances = shortest_distance_acyclic(&fst).unwrap();

        assert_eq!(*distances[0].value(), 0.0); // Start state
        assert_eq!(*distances[1].value(), 1.0); // Via s0 -> s1
        assert_eq!(*distances[2].value(), 1.5); // min(3.0, 1.0 + 0.5) = 1.5
    }

    #[test]
    fn test_acyclic_shortest_distance_chain() {
        let mut fst = VectorFst::<TropicalWeight>::new();
        let mut states = Vec::new();

        for _ in 0..5 {
            states.push(fst.add_state());
        }

        fst.set_start(states[0]);
        fst.set_final(states[4], TropicalWeight::one());

        for i in 0..4 {
            fst.add_arc(
                states[i],
                crate::arc::Arc::new(1, 1, TropicalWeight::new(1.0), states[i + 1]),
            );
        }

        let distances = shortest_distance_acyclic(&fst).unwrap();

        for (i, dist) in distances.iter().enumerate() {
            assert_eq!(*dist.value(), i as f32);
        }
    }

    #[test]
    fn test_acyclic_shortest_distance_unreachable() {
        let mut fst = VectorFst::<TropicalWeight>::new();
        let s0 = fst.add_state();
        let s1 = fst.add_state();
        let _s2 = fst.add_state(); // Unreachable

        fst.set_start(s0);
        fst.set_final(s1, TropicalWeight::one());
        fst.add_arc(s0, crate::arc::Arc::new(1, 1, TropicalWeight::new(1.0), s1));

        let distances = shortest_distance_acyclic(&fst).unwrap();

        assert_eq!(*distances[0].value(), 0.0);
        assert_eq!(*distances[1].value(), 1.0);
        assert!(Semiring::is_zero(&distances[2])); // Unreachable
    }

    #[test]
    fn test_cyclic_fst_detection() {
        let mut fst = VectorFst::<TropicalWeight>::new();
        let s0 = fst.add_state();
        let s1 = fst.add_state();

        fst.set_start(s0);
        fst.set_final(s1, TropicalWeight::one());

        // Create a cycle
        fst.add_arc(s0, crate::arc::Arc::new(1, 1, TropicalWeight::new(1.0), s1));
        fst.add_arc(s1, crate::arc::Arc::new(2, 2, TropicalWeight::new(1.0), s0));

        assert!(shortest_distance_acyclic::<TropicalWeight, _>(&fst).is_err());
        assert!(!is_acyclic::<TropicalWeight, _>(&fst));
    }

    #[test]
    fn test_is_acyclic() {
        // Acyclic FST
        let mut fst1 = VectorFst::<TropicalWeight>::new();
        let s0 = fst1.add_state();
        let s1 = fst1.add_state();
        fst1.set_start(s0);
        fst1.set_final(s1, TropicalWeight::one());
        fst1.add_arc(s0, crate::arc::Arc::new(1, 1, TropicalWeight::one(), s1));

        assert!(is_acyclic(&fst1));

        // Cyclic FST
        let mut fst2 = VectorFst::<TropicalWeight>::new();
        let t0 = fst2.add_state();
        let t1 = fst2.add_state();
        fst2.set_start(t0);
        fst2.add_arc(t0, crate::arc::Arc::new(1, 1, TropicalWeight::one(), t1));
        fst2.add_arc(t1, crate::arc::Arc::new(2, 2, TropicalWeight::one(), t0));

        assert!(!is_acyclic(&fst2));
    }

    #[test]
    fn test_reverse_shortest_distance() {
        let mut fst = VectorFst::<TropicalWeight>::new();
        let s0 = fst.add_state();
        let s1 = fst.add_state();
        let s2 = fst.add_state();

        fst.set_start(s0);
        fst.set_final(s2, TropicalWeight::new(0.5));

        fst.add_arc(s0, crate::arc::Arc::new(1, 1, TropicalWeight::new(1.0), s1));
        fst.add_arc(s1, crate::arc::Arc::new(2, 2, TropicalWeight::new(2.0), s2));

        let distances = shortest_distance_acyclic_reverse(&fst).unwrap();

        // s2 -> final: 0.5 (final weight)
        assert_eq!(*distances[2].value(), 0.5);
        // s1 -> s2 -> final: 2.0 + 0.5 = 2.5
        assert_eq!(*distances[1].value(), 2.5);
        // s0 -> s1 -> s2 -> final: 1.0 + 2.0 + 0.5 = 3.5
        assert_eq!(*distances[0].value(), 3.5);
    }

    #[test]
    fn test_empty_fst() {
        let fst = VectorFst::<TropicalWeight>::new();
        let distances = shortest_distance_acyclic(&fst).unwrap();
        assert!(distances.is_empty());
    }

    #[test]
    fn test_auto_selects_acyclic() {
        let mut fst = VectorFst::<TropicalWeight>::new();
        let s0 = fst.add_state();
        let s1 = fst.add_state();
        fst.set_start(s0);
        fst.set_final(s1, TropicalWeight::one());
        fst.add_arc(s0, crate::arc::Arc::new(1, 1, TropicalWeight::new(1.0), s1));

        // Auto should work for acyclic FSTs
        let distances = shortest_distance_auto(&fst).unwrap();
        assert_eq!(*distances[1].value(), 1.0);
    }
}