arcweight 0.3.0

//! Path iteration utilities for FSTs.
//!
//! This module provides iterators for enumerating all accepting paths through
//! a weighted finite-state transducer, with support for filtering, cycle
//! detection, and symbol table integration.
//!
//! # Overview
//!
//! Path iteration is useful for:
//! - Extracting all input/output string pairs from an FST
//! - Finding the n-best paths by weight
//! - Debugging FST structure by examining accepted sequences
//! - Converting FSTs to explicit path representations
//!
//! # Cycle Handling
//!
//! The iterators automatically detect and skip cycles to ensure termination.
//! For cyclic FSTs, only acyclic paths are enumerated.
//!
//! # Examples
//!
//! ## Enumerate All Paths
//!
//! ```
//! use arcweight::prelude::*;
//! use arcweight::utils::PathIterExt;
//!
//! 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(0.5), s1));
//! fst.add_arc(s0, Arc::new(2, 2, TropicalWeight::new(1.0), s1));
//! fst.add_arc(s1, Arc::new(3, 3, TropicalWeight::new(0.5), s2));
//!
//! for path in fst.paths_iter() {
//!     println!("Input: {:?}, Weight: {:?}", path.input_labels(), path.weight);
//! }
//! ```
//!
//! ## Limit Number of Paths
//!
//! ```
//! use arcweight::prelude::*;
//! use arcweight::utils::PathIterExt;
//!
//! let mut fst = VectorFst::<TropicalWeight>::new();
//! let s0 = fst.add_state();
//! fst.set_start(s0);
//! fst.set_final(s0, TropicalWeight::one());
//!
//! // Get at most 10 paths
//! let paths: Vec<_> = fst.paths_iter().with_max_paths(10).collect();
//! ```
//!
//! # Complexity
//!
//! - **Time**: O(number of paths) in the acyclic case
//! - **Space**: O(path length) for the current path being explored
//!
//! # References
//!
//! - Mehryar Mohri and Michael Riley. 2002. An efficient algorithm for the
//!   n-best-strings problem. In *Proc. ICSLP 2002*, 1313-1316.

use crate::arc::Arc;
use crate::fst::{Fst, StateId};
use crate::semiring::Semiring;
use std::collections::VecDeque;

/// Represents a complete accepting path through an FST.
///
/// An `FstPath` contains all information about a path from the start state
/// to a final state, including the sequence of states, arcs, and the total
/// path weight (product of arc weights times final weight).
///
/// # Examples
///
/// ```
/// use arcweight::prelude::*;
/// use arcweight::utils::PathIterExt;
///
/// 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, 2, TropicalWeight::new(0.5), s1));
///
/// let path = fst.paths_iter().next().unwrap();
/// assert_eq!(path.input_labels(), vec![1]);
/// assert_eq!(path.output_labels(), vec![2]);
/// assert_eq!(path.states, vec![s0, s1]);
/// ```
#[derive(Debug, Clone, PartialEq)]
pub struct FstPath<W: Semiring> {
    /// Sequence of states visited, starting from start state.
    pub states: Vec<StateId>,
    /// Sequence of arcs traversed along the path.
    pub arcs: Vec<Arc<W>>,
    /// Total weight of the path (product of arc weights and final weight).
    pub weight: W,
    /// The final state where the path terminates.
    pub final_state: StateId,
}

impl<W: Semiring> FstPath<W> {
    /// Extracts the sequence of input labels from the path.
    ///
    /// # Returns
    ///
    /// A vector of input labels in traversal order.
    pub fn input_labels(&self) -> Vec<u32> {
        self.arcs.iter().map(|arc| arc.ilabel).collect()
    }

    /// Extracts the sequence of output labels from the path.
    ///
    /// # Returns
    ///
    /// A vector of output labels in traversal order.
    pub fn output_labels(&self) -> Vec<u32> {
        self.arcs.iter().map(|arc| arc.olabel).collect()
    }

    /// Extracts input/output label pairs from the path.
    ///
    /// # Returns
    ///
    /// A vector of (input_label, output_label) tuples in traversal order.
    pub fn io_pairs(&self) -> Vec<(u32, u32)> {
        self.arcs
            .iter()
            .map(|arc| (arc.ilabel, arc.olabel))
            .collect()
    }
}

/// Iterator over all accepting paths in an FST.
///
/// This iterator performs a breadth-first search from the start state,
/// yielding each complete path to a final state. Cycles are automatically
/// detected and skipped to ensure termination.
///
/// # Complexity
///
/// - **Time**: O(number of acyclic paths)
/// - **Space**: O(maximum path length) for the exploration queue
///
/// # Examples
///
/// ```
/// use arcweight::prelude::*;
/// use arcweight::utils::{PathsIterator, PathIterExt};
///
/// 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(0.5), s1));
///
/// let paths: Vec<_> = fst.paths_iter().collect();
/// assert_eq!(paths.len(), 1);
/// ```
#[derive(Debug)]
pub struct PathsIterator<'a, W: Semiring, F: Fst<W>> {
    fst: &'a F,
    queue: VecDeque<(Vec<StateId>, Vec<Arc<W>>, W, StateId)>,
    max_paths: Option<usize>,
    paths_found: usize,
    weight_threshold: Option<W>,
}

impl<'a, W: Semiring, F: Fst<W>> PathsIterator<'a, W, F> {
    /// Creates a new path iterator for the given FST.
    ///
    /// # Arguments
    ///
    /// * `fst` - The FST to iterate over
    pub fn new(fst: &'a F) -> Self {
        let mut queue = VecDeque::new();
        if let Some(start) = fst.start() {
            queue.push_back((vec![start], Vec::new(), W::one(), start));
        }
        Self {
            fst,
            queue,
            max_paths: None,
            paths_found: 0,
            weight_threshold: None,
        }
    }

    /// Sets the maximum number of paths to return.
    ///
    /// After returning this many paths, the iterator will stop even if more
    /// paths exist. Useful for n-best path extraction.
    ///
    /// # Arguments
    ///
    /// * `max` - Maximum number of paths to yield
    ///
    /// # Examples
    ///
    /// ```
    /// use arcweight::prelude::*;
    /// use arcweight::utils::PathIterExt;
    ///
    /// let mut fst = VectorFst::<TropicalWeight>::new();
    /// let s0 = fst.add_state();
    /// fst.set_start(s0);
    /// fst.set_final(s0, TropicalWeight::one());
    ///
    /// // Get at most 5 paths
    /// let paths: Vec<_> = fst.paths_iter().with_max_paths(5).collect();
    /// assert!(paths.len() <= 5);
    /// ```
    pub fn with_max_paths(mut self, max: usize) -> Self {
        self.max_paths = Some(max);
        self
    }

    /// Sets a weight threshold for path filtering.
    ///
    /// Only paths with weight less than or equal to the threshold will be
    /// returned. For tropical semiring, this filters by path cost.
    ///
    /// # Type Requirements
    ///
    /// Requires a [`NaturallyOrderedSemiring`](crate::semiring::NaturallyOrderedSemiring)
    /// for weight comparison.
    ///
    /// # Arguments
    ///
    /// * `threshold` - Maximum weight for returned paths
    pub fn with_weight_threshold(mut self, threshold: W) -> Self
    where
        W: crate::semiring::NaturallyOrderedSemiring,
    {
        self.weight_threshold = Some(threshold);
        self
    }
}

impl<'a, W: Semiring, F: Fst<W>> Iterator for PathsIterator<'a, W, F>
where
    W: Clone,
{
    type Item = FstPath<W>;

    fn next(&mut self) -> Option<Self::Item> {
        while let Some((states, arcs, weight, current_state)) = self.queue.pop_front() {
            // Check if we've exceeded max paths
            if let Some(max) = self.max_paths {
                if self.paths_found >= max {
                    return None;
                }
            }

            // Check weight threshold (only if set and semiring is naturally ordered)
            if let Some(ref threshold) = self.weight_threshold {
                // This comparison is safe because with_weight_threshold requires NaturallyOrderedSemiring
                if weight > *threshold {
                    continue;
                }
            }

            // Check if current state is final
            if let Some(final_weight) = self.fst.final_weight(current_state) {
                let total_weight = weight.clone() * final_weight.clone();

                // Check final weight threshold
                if let Some(ref threshold) = self.weight_threshold {
                    if total_weight > *threshold {
                        // Continue exploring, but don't return this path
                    } else {
                        self.paths_found += 1;
                        return Some(FstPath {
                            states: states.clone(),
                            arcs: arcs.clone(),
                            weight: total_weight,
                            final_state: current_state,
                        });
                    }
                } else {
                    self.paths_found += 1;
                    return Some(FstPath {
                        states: states.clone(),
                        arcs: arcs.clone(),
                        weight: total_weight,
                        final_state: current_state,
                    });
                }
            }

            // Explore outgoing arcs
            for arc in self.fst.arcs(current_state) {
                // Check for cycles - if the next state is already in the current path, skip it
                // This ensures we detect cycles regardless of path length
                if states.contains(&arc.nextstate) {
                    continue;
                }

                let next_weight = weight.clone() * arc.weight.clone();
                let mut next_states = states.clone();
                next_states.push(arc.nextstate);
                let mut next_arcs = arcs.clone();
                next_arcs.push(arc.clone());

                self.queue
                    .push_back((next_states, next_arcs, next_weight, arc.nextstate));
            }
        }

        None
    }
}

/// Iterator over paths with human-readable string representations.
///
/// This iterator wraps [`PathsIterator`] and uses symbol tables to convert
/// numeric labels to human-readable strings.
///
/// # Examples
///
/// ```
/// use arcweight::prelude::*;
/// use arcweight::utils::{PathIterExt, SymbolTable};
///
/// let mut fst = VectorFst::<TropicalWeight>::new();
/// let mut syms = SymbolTable::new();
/// let hello = syms.add_symbol("hello");
///
/// 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(hello, hello, TropicalWeight::one(), s1));
///
/// for path in fst.string_paths_iter(Some(&syms), Some(&syms)) {
///     println!("Input: {}, Output: {}", path.input, path.output);
/// }
/// ```
#[derive(Debug)]
pub struct StringPathsIterator<'a, W: Semiring, F: Fst<W>> {
    path_iter: PathsIterator<'a, W, F>,
    input_symbols: Option<&'a crate::utils::SymbolTable>,
    output_symbols: Option<&'a crate::utils::SymbolTable>,
}

impl<'a, W: Semiring, F: Fst<W>> StringPathsIterator<'a, W, F> {
    /// Create a new string paths iterator
    pub fn new(
        fst: &'a F,
        input_symbols: Option<&'a crate::utils::SymbolTable>,
        output_symbols: Option<&'a crate::utils::SymbolTable>,
    ) -> Self {
        Self {
            path_iter: PathsIterator::new(fst),
            input_symbols,
            output_symbols,
        }
    }

    /// Set maximum number of paths
    pub fn with_max_paths(mut self, max: usize) -> Self {
        self.path_iter = self.path_iter.with_max_paths(max);
        self
    }
}

/// Human-readable string representation of a path.
///
/// Contains the input and output sequences as space-separated strings,
/// and the path weight as a string.
#[derive(Debug, Clone)]
pub struct StringPath {
    /// Input string (space-separated symbols from symbol table).
    pub input: String,
    /// Output string (space-separated symbols from symbol table).
    pub output: String,
    /// Path weight formatted as a string.
    pub weight: String,
}

impl<'a, W: Semiring, F: Fst<W>> Iterator for StringPathsIterator<'a, W, F>
where
    W: Clone + std::fmt::Display,
{
    type Item = StringPath;

    fn next(&mut self) -> Option<Self::Item> {
        let path = self.path_iter.next()?;

        let input = if let Some(symbols) = self.input_symbols {
            let labels: Vec<&str> = path
                .input_labels()
                .iter()
                .filter_map(|&label| symbols.find(label))
                .collect();
            labels.join(" ")
        } else {
            let labels: Vec<String> = path.input_labels().iter().map(|&l| l.to_string()).collect();
            labels.join(" ")
        };

        let output = if let Some(symbols) = self.output_symbols {
            let labels: Vec<&str> = path
                .output_labels()
                .iter()
                .filter_map(|&label| symbols.find(label))
                .collect();
            labels.join(" ")
        } else {
            let labels: Vec<String> = path
                .output_labels()
                .iter()
                .map(|&l| l.to_string())
                .collect();
            labels.join(" ")
        };

        Some(StringPath {
            input,
            output,
            weight: path.weight.to_string(),
        })
    }
}

/// Extension trait for FSTs to provide path iteration.
///
/// This trait is automatically implemented for all types that implement
/// [`Fst`], providing convenient path iteration methods.
///
/// # Examples
///
/// ```
/// use arcweight::prelude::*;
/// use arcweight::utils::PathIterExt;
///
/// let mut fst = VectorFst::<TropicalWeight>::new();
/// let s0 = fst.add_state();
/// fst.set_start(s0);
/// fst.set_final(s0, TropicalWeight::one());
///
/// // Numeric labels
/// for path in fst.paths_iter() {
///     println!("Weight: {:?}", path.weight);
/// }
///
/// // String labels (requires symbol tables)
/// for path in fst.string_paths_iter(None, None) {
///     println!("{} -> {} ({})", path.input, path.output, path.weight);
/// }
/// ```
pub trait PathIterExt<W: Semiring>: crate::fst::Fst<W> {
    /// Returns an iterator over all accepting paths in the FST.
    ///
    /// Paths are yielded in breadth-first order. Cycles are automatically
    /// detected and skipped.
    fn paths_iter(&self) -> PathsIterator<'_, W, Self>
    where
        Self: Sized;

    /// Returns an iterator over paths with human-readable string representations.
    ///
    /// # Arguments
    ///
    /// * `input_symbols` - Optional symbol table for input labels
    /// * `output_symbols` - Optional symbol table for output labels
    fn string_paths_iter<'a>(
        &'a self,
        input_symbols: Option<&'a crate::utils::SymbolTable>,
        output_symbols: Option<&'a crate::utils::SymbolTable>,
    ) -> StringPathsIterator<'a, W, Self>
    where
        Self: Sized;
}

impl<W: Semiring, F: Fst<W>> PathIterExt<W> for F {
    fn paths_iter(&self) -> PathsIterator<'_, W, Self> {
        PathsIterator::new(self)
    }

    fn string_paths_iter<'a>(
        &'a self,
        input_symbols: Option<&'a crate::utils::SymbolTable>,
        output_symbols: Option<&'a crate::utils::SymbolTable>,
    ) -> StringPathsIterator<'a, W, Self> {
        StringPathsIterator::new(self, input_symbols, output_symbols)
    }
}

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

    #[test]
    fn test_fst_path_input_labels() {
        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, 10, TropicalWeight::one(), s1));

        let path = fst.paths_iter().next().unwrap();
        assert_eq!(path.input_labels(), vec![1]);
    }

    #[test]
    fn test_fst_path_output_labels() {
        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, 10, TropicalWeight::one(), s1));

        let path = fst.paths_iter().next().unwrap();
        assert_eq!(path.output_labels(), vec![10]);
    }

    #[test]
    fn test_fst_path_io_pairs() {
        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, 10, TropicalWeight::one(), s1));

        let path = fst.paths_iter().next().unwrap();
        assert_eq!(path.io_pairs(), vec![(1, 10)]);
    }

    #[test]
    fn test_paths_iterator_empty() {
        let fst = VectorFst::<TropicalWeight>::new();
        let mut iter = fst.paths_iter();
        assert!(iter.next().is_none());
    }

    #[test]
    fn test_paths_iterator_single_path() {
        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::one(), s1));

        let paths: Vec<_> = fst.paths_iter().collect();
        assert_eq!(paths.len(), 1);
    }

    #[test]
    fn test_string_paths_iterator() {
        let mut fst = VectorFst::<TropicalWeight>::new();
        let mut symbols = SymbolTable::new();
        let hello_id = symbols.add_symbol("hello");

        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(hello_id, hello_id, TropicalWeight::one(), s1));

        let paths: Vec<_> = fst
            .string_paths_iter(Some(&symbols), Some(&symbols))
            .collect();
        assert_eq!(paths.len(), 1);
        assert_eq!(paths[0].input, "hello");
    }
}