arcweight 0.3.0

//! Topological sorting algorithm for acyclic FSTs.
//!
//! Reorders states in a directed acyclic graph (DAG) FST so that for every arc $(u,v)$,
//! state $u$ appears before state $v$ in the ordering. This enables efficient dynamic
//! programming algorithms that process states in dependency order.
//!
//! Topological sorting is essential for algorithms like shortest-distance computation on
//! acyclic graphs, where processing states in topological order guarantees that all
//! predecessors are processed before their successors.
//!
//! # Complexity
//!
//! - **Time:** $`O(|V| + |E|)`$ using depth-first search
//! - **Space:** $`O(|V|)`$ for visited sets and the result ordering
//!
//! # References
//!
//! - Thomas H. Cormen, Charles E. Leiserson, Ronald L. Rivest, and Clifford Stein.
//!   2009. *Introduction to Algorithms* (3rd ed.). MIT Press, Cambridge, MA.
//!   Chapter 22: Elementary Graph Algorithms.
//! - Mehryar Mohri. 2009. Weighted automata algorithms. In *Handbook of Weighted
//!   Automata*, Manfred Droste, Werner Kuich, and Heiko Vogler (Eds.). Springer,
//!   Berlin, Heidelberg, 213-254.

use crate::arc::Arc;
use crate::fst::{Fst, MutableFst, StateId};
use crate::semiring::Semiring;
use crate::{Error, Result};
use std::collections::HashSet;

/// Topologically sorts states in an acyclic FST for optimal processing order.
///
/// Reorders the states of a directed acyclic graph (DAG) FST so that all arcs
/// point "forward" in the state ordering. For every arc $(u, v)$, state $u$
/// will appear before state $v$ in the result.
///
/// This enables efficient dynamic programming algorithms that process states
/// in dependency order, ensuring each state is processed after all its predecessors.
///
/// # Type Parameters
///
/// * `W` - Weight type implementing [`Semiring`]
/// * `F` - Input FST type implementing [`Fst<W>`]
/// * `M` - Output FST type implementing [`MutableFst<W>`] and [`Default`]
///
/// # Returns
///
/// A new FST with states renumbered in topological order.
///
/// # Errors
///
/// Returns [`Error::Algorithm`] if:
/// - The input FST contains cycles (not a DAG)
/// - The FST structure is invalid
///
/// # Examples
///
/// ```rust
/// use arcweight::prelude::*;
///
/// // Create DAG FST: 0 -> 1 -> 2, 0 -> 2
/// 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());
///
/// fst.add_arc(s0, Arc::new('a' as u32, 'a' as u32, TropicalWeight::one(), s1));
/// fst.add_arc(s1, Arc::new('b' as u32, 'b' as u32, TropicalWeight::one(), s2));
/// fst.add_arc(s0, Arc::new('c' as u32, 'c' as u32, TropicalWeight::one(), s2));
///
/// // Sort states in topological order
/// let sorted: VectorFst<TropicalWeight> = topsort(&fst)?;
///
/// // All arcs now point to higher-numbered states
/// for state in sorted.states() {
///     for arc in sorted.arcs(state) {
///         assert!(state < arc.nextstate);
///     }
/// }
/// # Ok::<(), Box<dyn std::error::Error>>(())
/// ```
///
/// # References
///
/// - Thomas H. Cormen, Charles E. Leiserson, Ronald L. Rivest, and Clifford Stein.
///   2009. *Introduction to Algorithms* (3rd ed.). MIT Press, Cambridge, MA.
///   Chapter 22: Elementary Graph Algorithms.
///
/// [`Semiring`]: crate::semiring::Semiring
pub fn topsort<W, F, M>(fst: &F) -> Result<M>
where
    W: Semiring,
    F: Fst<W>,
    M: MutableFst<W> + Default,
{
    // compute topological order
    let order = compute_topological_order(fst)?;

    // create mapping from old to new state IDs
    let mut state_map = vec![None; fst.num_states()];
    for (new_id, &old_id) in order.iter().enumerate() {
        state_map[old_id as usize] = Some(new_id as StateId);
    }

    let mut result = M::default();

    // create states in topological order
    for _ in &order {
        result.add_state();
    }

    // set start
    if let Some(start) = fst.start() {
        if let Some(new_start) = state_map[start as usize] {
            result.set_start(new_start);
        }
    }

    // copy with remapped states
    for &old_state in &order {
        if let Some(new_state) = state_map[old_state as usize] {
            // final weight
            if let Some(weight) = fst.final_weight(old_state) {
                result.set_final(new_state, weight.clone());
            }

            // arcs
            for arc in fst.arcs(old_state) {
                if let Some(new_nextstate) = state_map[arc.nextstate as usize] {
                    result.add_arc(
                        new_state,
                        Arc::new(arc.ilabel, arc.olabel, arc.weight.clone(), new_nextstate),
                    );
                }
            }
        }
    }

    Ok(result)
}

fn compute_topological_order<W: Semiring, F: Fst<W>>(fst: &F) -> Result<Vec<StateId>> {
    let mut visited = HashSet::new();
    let mut finished = HashSet::new();
    let mut order = Vec::new();

    fn dfs<W: Semiring, F: Fst<W>>(
        fst: &F,
        state: StateId,
        visited: &mut HashSet<StateId>,
        finished: &mut HashSet<StateId>,
        order: &mut Vec<StateId>,
    ) -> Result<()> {
        visited.insert(state);

        for arc in fst.arcs(state) {
            if !visited.contains(&arc.nextstate) {
                dfs(fst, arc.nextstate, visited, finished, order)?;
            } else if !finished.contains(&arc.nextstate) {
                return Err(Error::Algorithm("FST has cycles".into()));
            }
        }

        finished.insert(state);
        order.push(state);
        Ok(())
    }

    // start DFS from start state if it exists
    if let Some(start) = fst.start() {
        if !visited.contains(&start) {
            dfs(fst, start, &mut visited, &mut finished, &mut order)?;
        }
    }

    // visit any remaining unvisited states
    for state in fst.states() {
        if !visited.contains(&state) {
            dfs(fst, state, &mut visited, &mut finished, &mut order)?;
        }
    }

    order.reverse();
    Ok(order)
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::prelude::*;
    use num_traits::One;

    #[test]
    fn test_topsort() {
        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());

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

        let sorted: VectorFst<TropicalWeight> = topsort(&fst).unwrap();

        // Verify topological order: all arcs go from lower to higher state IDs
        for state in sorted.states() {
            for arc in sorted.arcs(state) {
                assert!(
                    state < arc.nextstate,
                    "Arc from {} to {} violates topological order",
                    state,
                    arc.nextstate
                );
            }
        }

        assert!(sorted.start().is_some());
        assert_eq!(sorted.num_states(), fst.num_states());
    }

    #[test]
    fn test_topsort_cyclic() {
        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, Arc::new(1, 1, TropicalWeight::new(1.0), s1));
        fst.add_arc(s1, Arc::new(2, 2, TropicalWeight::new(1.0), s0));

        // Topsort should fail or handle cycles appropriately
        let result =
            topsort::<TropicalWeight, VectorFst<TropicalWeight>, VectorFst<TropicalWeight>>(&fst);
        if let Ok(sorted) = result {
            // If it succeeds, should still have valid structure
            assert!(sorted.start().is_some());
        }
        // If it fails, that's also acceptable for cyclic graphs
    }

    #[test]
    fn test_topsort_empty_fst() {
        let fst = VectorFst::<TropicalWeight>::new();
        let sorted: VectorFst<TropicalWeight> = topsort(&fst).unwrap();

        assert_eq!(sorted.num_states(), 0);
        assert!(sorted.is_empty());
    }

    #[test]
    fn test_topsort_single_state() {
        let mut fst = VectorFst::<TropicalWeight>::new();
        let s0 = fst.add_state();
        fst.set_start(s0);
        fst.set_final(s0, TropicalWeight::new(2.0));

        let sorted: VectorFst<TropicalWeight> = topsort(&fst).unwrap();

        assert_eq!(sorted.num_states(), 1);
        assert_eq!(sorted.start(), Some(0));
        assert!(sorted.is_final(0));
        assert_eq!(sorted.final_weight(0), Some(&TropicalWeight::new(2.0)));
    }

    #[test]
    fn test_topsort_linear_chain() {
        let mut fst = VectorFst::<TropicalWeight>::new();
        let states: Vec<_> = (0..5).map(|_| fst.add_state()).collect();

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

        // Create linear chain: 0 -> 1 -> 2 -> 3 -> 4
        for i in 0..4 {
            fst.add_arc(
                states[i],
                Arc::new(
                    (i + 1) as u32,
                    (i + 1) as u32,
                    TropicalWeight::new(i as f32),
                    states[i + 1],
                ),
            );
        }

        let sorted: VectorFst<TropicalWeight> = topsort(&fst).unwrap();

        // Verify topological ordering is maintained
        assert_eq!(sorted.num_states(), fst.num_states());
        assert!(sorted.start().is_some());

        // All arcs should go forward in the sorted order
        for state in sorted.states() {
            for arc in sorted.arcs(state) {
                assert!(state < arc.nextstate);
            }
        }
    }

    #[test]
    fn test_topsort_diamond_dag() {
        let mut fst = VectorFst::<TropicalWeight>::new();
        let s0 = fst.add_state(); // Source
        let s1 = fst.add_state(); // Left branch
        let s2 = fst.add_state(); // Right branch
        let s3 = fst.add_state(); // Sink

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

        // Diamond pattern: 0 -> {1,2} -> 3
        fst.add_arc(s0, Arc::new(1, 1, TropicalWeight::one(), s1));
        fst.add_arc(s0, Arc::new(2, 2, TropicalWeight::one(), s2));
        fst.add_arc(s1, Arc::new(3, 3, TropicalWeight::one(), s3));
        fst.add_arc(s2, Arc::new(4, 4, TropicalWeight::one(), s3));

        let sorted: VectorFst<TropicalWeight> = topsort(&fst).unwrap();

        assert_eq!(sorted.num_states(), 4);
        assert!(sorted.start().is_some());

        // Verify topological ordering
        for state in sorted.states() {
            for arc in sorted.arcs(state) {
                assert!(state < arc.nextstate);
            }
        }
    }

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

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

        // Connected component: 0 -> 1
        fst.add_arc(s0, Arc::new(1, 1, TropicalWeight::one(), s1));

        // Disconnected component: 2 -> 3
        fst.add_arc(s2, Arc::new(2, 2, TropicalWeight::one(), s3));

        let sorted: VectorFst<TropicalWeight> = topsort(&fst).unwrap();

        assert_eq!(sorted.num_states(), 4);
        assert!(sorted.start().is_some());

        // All arcs should still maintain topological order
        for state in sorted.states() {
            for arc in sorted.arcs(state) {
                assert!(state < arc.nextstate);
            }
        }
    }

    #[test]
    fn test_topsort_preserves_weights() {
        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(3.5));

        fst.add_arc(s0, Arc::new(1, 1, TropicalWeight::new(1.2), s1));
        fst.add_arc(s1, Arc::new(2, 2, TropicalWeight::new(2.3), s2));

        let sorted: VectorFst<TropicalWeight> = topsort(&fst).unwrap();

        // Weights should be preserved
        assert_eq!(sorted.num_states(), fst.num_states());

        // Check that weights are preserved (structure may change but semantics preserved)
        let total_weight_orig: f32 = fst
            .states()
            .flat_map(|s| fst.arcs(s))
            .map(|arc| *arc.weight.value())
            .sum();
        let total_weight_sorted: f32 = sorted
            .states()
            .flat_map(|s| sorted.arcs(s))
            .map(|arc| *arc.weight.value())
            .sum();

        assert!((total_weight_orig - total_weight_sorted).abs() < 1e-6);
    }

    #[test]
    fn test_compute_topological_order() {
        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.add_arc(s0, Arc::new(1, 1, TropicalWeight::one(), s1));
        fst.add_arc(s1, Arc::new(2, 2, TropicalWeight::one(), s2));

        let order = compute_topological_order(&fst).unwrap();

        assert_eq!(order.len(), 3);

        // Should contain all states
        assert!(order.contains(&s0));
        assert!(order.contains(&s1));
        assert!(order.contains(&s2));

        // Find positions in the ordering
        let pos0 = order.iter().position(|&x| x == s0).unwrap();
        let pos1 = order.iter().position(|&x| x == s1).unwrap();
        let pos2 = order.iter().position(|&x| x == s2).unwrap();

        // Should respect dependencies: s0 before s1 before s2
        assert!(pos0 < pos1);
        assert!(pos1 < pos2);
    }

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

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

        // Self-loop creates a cycle
        fst.add_arc(s0, Arc::new(1, 1, TropicalWeight::one(), s0));

        // Should detect the cycle
        let result = compute_topological_order(&fst);
        assert!(result.is_err());
    }

    #[test]
    fn test_topsort_complex_dag() {
        let mut fst = VectorFst::<TropicalWeight>::new();
        let states: Vec<_> = (0..6).map(|_| fst.add_state()).collect();

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

        // Complex DAG structure
        fst.add_arc(states[0], Arc::new(1, 1, TropicalWeight::one(), states[1]));
        fst.add_arc(states[0], Arc::new(2, 2, TropicalWeight::one(), states[2]));
        fst.add_arc(states[1], Arc::new(3, 3, TropicalWeight::one(), states[3]));
        fst.add_arc(states[2], Arc::new(4, 4, TropicalWeight::one(), states[3]));
        fst.add_arc(states[1], Arc::new(5, 5, TropicalWeight::one(), states[4]));
        fst.add_arc(states[3], Arc::new(6, 6, TropicalWeight::one(), states[5]));
        fst.add_arc(states[4], Arc::new(7, 7, TropicalWeight::one(), states[5]));

        let sorted: VectorFst<TropicalWeight> = topsort(&fst).unwrap();

        assert_eq!(sorted.num_states(), 6);
        assert!(sorted.start().is_some());

        // Verify all arcs respect topological order
        for state in sorted.states() {
            for arc in sorted.arcs(state) {
                assert!(
                    state < arc.nextstate,
                    "Arc from {} to {} violates topological order",
                    state,
                    arc.nextstate
                );
            }
        }
    }
}