arcweight 0.3.0

//! Composition of weighted finite-state transducers.
//!
//! This module implements the fundamental composition operation for WFSTs, which
//! combines two transducers by matching the output of the first against the input
//! of the second. For transducers $`T_1: \Sigma^* \to \Gamma^*`$ and
//! $`T_2: \Gamma^* \to \Delta^*`$, their composition $`T_1 \circ T_2: \Sigma^* \to \Delta^*`$
//! computes all valid transformations through both transducers.
//!
//! # Mathematical Definition
//!
//! For weighted transducers over a semiring $`(K, \oplus, \otimes, \bar{0}, \bar{1})`$:
//!
//! ```text
//! (T₁ ∘ T₂)(x, z) = ⊕_{y ∈ Γ*} T₁(x, y) ⊗ T₂(y, z)
//! ```
//!
//! where $`x \in \Sigma^*`$, $`y \in \Gamma^*`$, $`z \in \Delta^*`$.
//!
//! # Complexity
//!
//! | Variant | Time | Space |
//! |---------|------|-------|
//! | [`compose`] | $`O(V_1 V_2 E_1 E_2)`$ worst | $`O(V_1 V_2)`$ |
//! | [`compose_sorted`] | $`O(V_1 V_2 (E_1 \log E_1 + E_2 \log E_2))`$ | $`O(V_1 V_2)`$ |
//!
//! In practice, composition is often much faster due to label sparsity.
//!
//! # Algorithm
//!
//! The implementation uses on-the-fly (lazy) composition with epsilon handling:
//!
//! 1. Create start state from $`(q_{01}, q_{02})`$
//! 2. Explore reachable state pairs via BFS
//! 3. For each pair $`(q_1, q_2)`$, create arcs for matching labels
//! 4. Handle epsilon transitions with filter states to avoid redundant paths
//!
//! # Applications
//!
//! - **Speech recognition:** $`H \circ C \circ L \circ G`$ cascade
//! - **Machine translation:** Transfer model composition
//! - **Morphological analysis:** Lexicon + rules composition
//! - **Spell checking:** Error model + dictionary
//!
//! # References
//!
//! \[1\] Mohri, M., Pereira, F., and Riley, M. 2008. Speech recognition with weighted
//!     finite-state transducers. In *Springer Handbook of Speech Processing*, 559-584.
//!     <https://doi.org/10.1007/978-3-540-49127-9_28>
//!
//! \[2\] Pereira, F. C. N. and Riley, M. D. 1997. Speech recognition by composition
//!     of weighted finite automata. In *Finite-State Language Processing*, E. Roche
//!     and Y. Schabes, Eds. MIT Press, 431-453.
//!
//! \[3\] Mohri, M. 2009. Weighted automata algorithms. In *Handbook of Weighted
//!     Automata*, M. Droste, W. Kuich, and H. Vogler, Eds. Springer, 213-254.
//!     <https://doi.org/10.1007/978-3-642-01492-5_6>

use crate::arc::Arc;
use crate::fst::{Fst, MutableFst, StateId};
use crate::semiring::Semiring;
use crate::{Error, Result};
use rustc_hash::{FxHashMap, FxHashSet};
use std::collections::VecDeque;

/// Trait for composition filters that control arc matching behavior.
///
/// Composition filters enable customization of how arcs from two FSTs are matched
/// and combined during composition. The filter maintains state to handle epsilon
/// transitions and can implement various matching strategies.
///
/// # Filter State
///
/// The `FilterState` type tracks information needed for proper epsilon handling.
/// For simple label matching, `()` suffices. For epsilon-sequencing filters,
/// the state tracks which FST should advance on epsilon transitions.
///
/// # Implementing Custom Filters
///
/// ```rust
/// use arcweight::prelude::*;
///
/// struct MyFilter;
///
/// impl<W: Semiring> ComposeFilter<W> for MyFilter {
///     type FilterState = ();
///
///     fn start() -> Self::FilterState { () }
///
///     fn filter_arc(
///         &self,
///         arc1: &Arc<W>,
///         arc2: &Arc<W>,
///         _fs: &Self::FilterState,
///     ) -> Option<(Arc<W>, Self::FilterState)> {
///         // Custom matching logic
///         if arc1.olabel == arc2.ilabel {
///             Some((Arc::new(
///                 arc1.ilabel,
///                 arc2.olabel,
///                 arc1.weight.times(&arc2.weight),
///                 0,
///             ), ()))
///         } else {
///             None
///         }
///     }
/// }
/// ```
///
/// # See Also
///
/// - [`DefaultComposeFilter`] - Standard label-matching filter
/// - [`compose`] - Main composition function using filters
pub trait ComposeFilter<W: Semiring> {
    /// Associated state type for tracking filter context across transitions.
    ///
    /// For filters without state, use `()`. For epsilon-sequencing filters,
    /// use an enum to track which FST should advance.
    type FilterState: Clone + Default + Eq + std::hash::Hash;

    /// Returns the initial filter state for composition start.
    fn start() -> Self::FilterState;

    /// Determines if two arcs can be composed and returns the resulting arc.
    ///
    /// # Arguments
    ///
    /// * `arc1` - Arc from the first FST
    /// * `arc2` - Arc from the second FST
    /// * `fs` - Current filter state
    ///
    /// # Returns
    ///
    /// `Some((composed_arc, next_state))` if arcs can be composed, `None` otherwise.
    /// The composed arc's `nextstate` field is set by the caller.
    fn filter_arc(
        &self,
        arc1: &Arc<W>,
        arc2: &Arc<W>,
        fs: &Self::FilterState,
    ) -> Option<(Arc<W>, Self::FilterState)>;
}

/// Default composition filter implementing standard label matching.
///
/// This filter matches arcs when `arc1.olabel == arc2.ilabel`, handling epsilon
/// transitions by allowing them to match any label. The composed arc takes
/// `ilabel` from `arc1` and `olabel` from `arc2`.
///
/// # Epsilon Handling
///
/// - Epsilon output from FST1 (`arc1.olabel == 0`) matches any FST2 input
/// - Epsilon input in FST2 (`arc2.ilabel == 0`) matches any FST1 output
/// - Both epsilon: composed arc is epsilon
///
/// # Example
///
/// ```rust
/// use arcweight::prelude::*;
///
/// 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, Arc::new(1, 2, TropicalWeight::new(0.5), s1));
///
/// let mut fst2 = VectorFst::<TropicalWeight>::new();
/// let t0 = fst2.add_state();
/// let t1 = fst2.add_state();
/// fst2.set_start(t0);
/// fst2.set_final(t1, TropicalWeight::one());
/// fst2.add_arc(t0, Arc::new(2, 3, TropicalWeight::new(0.3), t1));
///
/// // Compose using default filter
/// let result: VectorFst<TropicalWeight> = compose_default(&fst1, &fst2)?;
/// # Ok::<(), arcweight::Error>(())
/// ```
#[derive(Debug, Clone, Copy, Default)]
pub struct DefaultComposeFilter;

impl<W: Semiring> ComposeFilter<W> for DefaultComposeFilter {
    type FilterState = ();

    fn start() -> Self::FilterState {}

    fn filter_arc(
        &self,
        arc1: &Arc<W>,
        arc2: &Arc<W>,
        _fs: &Self::FilterState,
    ) -> Option<(Arc<W>, Self::FilterState)> {
        use crate::fst::NO_LABEL;

        // Match when labels are equal (standard case)
        // Or when arc1 has epsilon output (matches any input in FST2)
        // Or when arc2 has epsilon input (matches any output from FST1)
        if arc1.olabel == arc2.ilabel || arc1.olabel == NO_LABEL || arc2.ilabel == NO_LABEL {
            // Determine composed arc labels
            // If both are epsilon: result is epsilon
            // Otherwise: input from arc1, output from arc2
            let composed_ilabel = if arc1.olabel == NO_LABEL && arc2.ilabel == NO_LABEL {
                NO_LABEL
            } else {
                arc1.ilabel
            };

            let composed_olabel = if arc1.olabel == NO_LABEL && arc2.ilabel == NO_LABEL {
                NO_LABEL
            } else {
                arc2.olabel
            };

            Some((
                Arc::new(
                    composed_ilabel,
                    composed_olabel,
                    arc1.weight.times(&arc2.weight),
                    0, // nextstate set later
                ),
                (),
            ))
        } else {
            None
        }
    }
}

/// Computes epsilon closure for composition with cycle detection.
///
/// Returns a map from reachable states to their accumulated epsilon weights.
/// Uses BFS with iteration limiting to handle cycles in non-idempotent semirings.
///
/// # Algorithm
///
/// 1. Initialize queue with `(start, W::one())`
/// 2. For each state, follow epsilon transitions accumulating weights
/// 3. Update closure using semiring plus when revisiting states
/// 4. Limit iterations per state to `MAX_ITERATIONS` for cycle handling
///
/// # Complexity
///
/// - **Time:** $`O(V \cdot k)`$ where $`k`$ is the iteration limit
/// - **Space:** $`O(V)`$ for closure map
///
/// # Arguments
///
/// * `fst` - Input FST to compute closure for
/// * `start` - State to compute closure from
///
/// # Returns
///
/// Map from each epsilon-reachable state to its accumulated weight.
fn compute_epsilon_closure_for_compose<W: Semiring, F: Fst<W>>(
    fst: &F,
    start: StateId,
) -> Result<FxHashMap<StateId, W>> {
    let mut closure = FxHashMap::default();
    let mut queue = VecDeque::new();
    let mut iteration_count: FxHashMap<StateId, usize> = FxHashMap::default();
    const MAX_ITERATIONS: usize = 1000; // Limit iterations to prevent infinite loops

    queue.push_back((start, W::one()));
    closure.insert(start, W::one());
    iteration_count.insert(start, 0);

    while let Some((state, weight)) = queue.pop_front() {
        // Skip if we've exceeded iteration limit (indicates potential non-convergent cycle)
        let iter = iteration_count.get(&state).copied().unwrap_or(0);
        if iter >= MAX_ITERATIONS {
            continue;
        }

        // Follow epsilon transitions
        for arc in fst.arcs(state) {
            if arc.is_epsilon() {
                let next_weight = weight.times(&arc.weight);
                let next_iter = iter + 1;

                let should_update = match closure.get(&arc.nextstate) {
                    Some(existing) => {
                        // Update if we found a better path (using semiring plus)
                        let combined = existing.plus(&next_weight);
                        if combined != *existing {
                            closure.insert(arc.nextstate, combined.clone());
                            true
                        } else {
                            false
                        }
                    }
                    None => {
                        closure.insert(arc.nextstate, next_weight.clone());
                        true
                    }
                };

                if should_update {
                    // Only add to queue if we haven't exceeded iteration limit
                    if next_iter < MAX_ITERATIONS {
                        iteration_count.insert(arc.nextstate, next_iter);
                        queue.push_back((arc.nextstate, closure[&arc.nextstate].clone()));
                    }
                }
            }
        }
    }

    Ok(closure)
}

/// Compose two FSTs using a custom composition filter
///
/// Computes the composition of two FSTs where the output of the first FST
/// is matched against the input of the second FST. The result accepts sequences
/// that are transformed by the first FST into sequences accepted by the second.
///
/// **Mathematical Definition:** For FSTs T₁: Σ* → Γ* and T₂: Γ* → Δ*,
/// their composition T₁ ∘ T₂: Σ* → Δ* computes:
/// ```text
/// (T₁ ∘ T₂)(x) = ⊕_{y ∈ Γ*} T₁(x,y) ⊗ T₂(y)
/// where:
///   x ∈ Σ* = input sequence
///   y ∈ Γ* = intermediate sequence
///   ⊕ = semiring addition (sum over all paths)
///   ⊗ = semiring multiplication (combine weights)
/// ```
///
/// # Complexity
///
/// - **Time:** O(V₁V₂E₁E₂) worst case, O(VE) average case
///   - V₁, V₂ = number of states in each FST
///   - E₁, E₂ = number of arcs in each FST
///   - Worst case: dense cross-product of all state-arc pairs
///   - Average case: sparse composition with few matching labels
/// - **Space:** O(V₁V₂) for reachable state pairs
///   - FxHashMap for state pair mapping
///   - Queue for BFS traversal
///
/// # Algorithm
///
/// On-the-fly composition using breadth-first state exploration:
/// 1. Create start state from (start₁, start₂, filter_start)
/// 2. While queue non-empty:
///    - Pop state triple (s₁, s₂, filter_state)
///    - If both s₁, s₂ are final: mark composed state as final with w₁ ⊗ w₂
///    - For each arc pair (a₁ from s₁, a₂ from s₂):
///      - Apply filter to check if arcs compose
///      - If match: create/reuse destination state, add composed arc
/// 3. Result contains only reachable states from start
///
/// Based on Mohri, Pereira, and Riley (2008) with lazy state construction.
///
/// # Performance Notes
///
/// - **On-the-fly construction:** Avoids creating unreachable state pairs
/// - **Label sparsity:** Performance improves with few matching labels between FSTs
/// - **Filter overhead:** Custom filters add per-arc-pair overhead
/// - **Memory pattern:** BFS traversal provides good cache locality
/// - **Optimization tip:** Pre-sort arcs by label for faster label matching
/// - **Best for:** Sparse FSTs with limited label overlap (e.g., NLP pipelines)
///
/// # Examples
///
/// ## Basic Composition
///
/// ```rust
/// use arcweight::prelude::*;
///
/// // Create first FST: maps 1 -> 2
/// 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, Arc::new(1, 2, TropicalWeight::new(0.5), s1));
///
/// // Create second FST: maps 2 -> 3
/// let mut fst2 = VectorFst::<TropicalWeight>::new();
/// let s0 = fst2.add_state();
/// let s1 = fst2.add_state();
/// fst2.set_start(s0);
/// fst2.set_final(s1, TropicalWeight::one());
/// fst2.add_arc(s0, Arc::new(2, 3, TropicalWeight::new(0.3), s1));
///
/// // Compose: result maps 1 -> 3 with weight 0.8
/// let filter = DefaultComposeFilter;
/// let composed: VectorFst<TropicalWeight> = compose(&fst1, &fst2, filter).unwrap();
/// assert!(composed.num_states() > 0);
/// ```
///
/// ## NLP Pipeline Example
///
/// ```rust
/// use arcweight::prelude::*;
///
/// // Example of chaining two FSTs for NLP pipeline
/// // FST 1: simple tokenization (example structure)
/// let mut tokenizer = VectorFst::<TropicalWeight>::new();
/// let s0 = tokenizer.add_state();
/// let s1 = tokenizer.add_state();
/// tokenizer.set_start(s0);
/// tokenizer.set_final(s1, TropicalWeight::one());
/// tokenizer.add_arc(s0, Arc::new(1, 2, TropicalWeight::one(), s1));
///
/// // FST 2: POS tagging (example structure)  
/// let mut pos_tagger = VectorFst::<TropicalWeight>::new();
/// let s0 = pos_tagger.add_state();
/// let s1 = pos_tagger.add_state();
/// pos_tagger.set_start(s0);
/// pos_tagger.set_final(s1, TropicalWeight::one());
/// pos_tagger.add_arc(s0, Arc::new(2, 3, TropicalWeight::one(), s1));
///
/// // Chain them: raw_text -> tagged_tokens
/// let filter = DefaultComposeFilter;
/// let pipeline: VectorFst<TropicalWeight> = compose(&tokenizer, &pos_tagger, filter).unwrap();
/// assert!(pipeline.num_states() > 0);
/// ```
///
/// # Custom Filters
///
/// Use custom composition filters for specialized epsilon handling or
/// special matching criteria:
///
/// ```rust
/// use arcweight::prelude::*;
///
/// struct CustomFilter {
///     allow_epsilon: bool,
/// }
///
/// impl<W: Semiring> ComposeFilter<W> for CustomFilter {
///     type FilterState = i32;
///     
///     fn start() -> Self::FilterState { 0 }
///     
///     fn filter_arc(
///         &self,
///         arc1: &Arc<W>,
///         arc2: &Arc<W>,
///         _fs: &Self::FilterState,
///     ) -> Option<(Arc<W>, Self::FilterState)> {
///         // Only compose if labels match
///         if arc1.olabel == arc2.ilabel {
///             Some((
///                 Arc::new(
///                     arc1.ilabel,
///                     arc2.olabel,
///                     arc1.weight.times(&arc2.weight),
///                     0, // nextstate set later
///                 ),
///                 0,
///             ))
///         } else {
///             None
///         }
///     }
/// }
///
/// // Use the custom filter
/// let mut fst1 = VectorFst::<TropicalWeight>::new();
/// let s0 = fst1.add_state();
/// fst1.set_start(s0);
/// fst1.set_final(s0, TropicalWeight::one());
///
/// let mut fst2 = VectorFst::<TropicalWeight>::new();
/// let s0 = fst2.add_state();
/// fst2.set_start(s0);
/// fst2.set_final(s0, TropicalWeight::one());
///
/// let filter = CustomFilter { allow_epsilon: true };
/// let result: VectorFst<TropicalWeight> = compose(&fst1, &fst2, filter).unwrap();
/// assert_eq!(result.num_states(), 1);
/// ```
///
/// # Errors
///
/// Returns [`Error::Algorithm`] if:
/// - Either input FST is invalid or has no start state
/// - Memory allocation fails during computation
/// - The composition filter produces invalid arc combinations
/// - State space becomes too large (implementation limit exceeded)
///
/// # References
///
/// \[1\] Mohri, M., Pereira, F., and Riley, M. 2008. Speech recognition with weighted
///     finite-state transducers. In *Springer Handbook of Speech Processing*, 559-584.
///     <https://doi.org/10.1007/978-3-540-49127-9_28>
///
/// \[2\] Pereira, F. C. N. and Riley, M. D. 1997. Speech recognition by composition
///     of weighted finite automata. In *Finite-State Language Processing*, 431-453.
///
/// # See Also
///
/// - [`compose_default`] - Convenience function using standard label matching
/// - [`ComposeFilter`] - Trait for implementing custom composition logic
/// - [`DefaultComposeFilter`] - Standard label-based composition filter
/// - [`determinize`] - Often applied after composition to reduce nondeterminism
/// - [`minimize`] - Reduce composed FST size
/// - [Fst trait](crate::fst::Fst) - Core FST interface
/// - [Semiring trait](crate::semiring::Semiring) - Weight algebra
///
/// [`determinize`]: crate::algorithms::determinize::determinize
/// [`minimize`]: crate::algorithms::minimize::minimize
pub fn compose<W, F1, F2, M, CF>(fst1: &F1, fst2: &F2, filter: CF) -> Result<M>
where
    W: Semiring,
    F1: Fst<W>,
    F2: Fst<W>,
    M: MutableFst<W> + Default,
    CF: ComposeFilter<W>,
{
    let start1 = fst1
        .start()
        .ok_or_else(|| Error::Algorithm("First FST has no start state".into()))?;
    let start2 = fst2
        .start()
        .ok_or_else(|| Error::Algorithm("Second FST has no start state".into()))?;

    let mut result = M::default();
    let mut state_map = FxHashMap::default();
    let mut queue = Vec::new();
    // Epsilon Handling Strategy:
    //
    // When processing a state pair (s1, s2), we compute epsilon closures for both FSTs.
    // The epsilon closure includes all states reachable via epsilon-only paths, with
    // accumulated weights. We then create arcs from ALL pairs in the Cartesian product
    // of the two closures.
    //
    // To avoid duplicate work:
    // - `processed_via_closure` tracks state pairs that have already had their arcs
    //   created as part of another state's epsilon closure processing.
    // - If (cs1, cs2) appears in the closure of (s1, s2) and we create arcs from it,
    //   we mark (cs1, cs2) as processed so when it appears in the queue later, we skip it.
    // - `created_epsilon_from_current` tracks whether we created epsilon arcs directly
    //   from the current state pair, which affects whether to mark certain pairs as processed.
    //
    // This avoids exponential blowup from repeatedly processing the same epsilon-reachable
    // state pairs while ensuring all valid paths are captured with correct weights.
    let mut processed_via_closure = FxHashSet::default();

    // create start state
    let start_state = result.add_state();
    result.set_start(start_state);
    state_map.insert((start1, start2, CF::start()), start_state);
    queue.push((start1, start2, CF::start(), start_state));

    // process states
    while let Some((s1, s2, fs, current)) = queue.pop() {
        // Skip if this state pair was already processed through epsilon closure
        if processed_via_closure.contains(&(s1, s2)) {
            continue;
        }
        // Compute epsilon closure for both FSTs
        // Epsilon closure for FST1 with cycle detection
        let closure1 = compute_epsilon_closure_for_compose(fst1, s1)?;

        // Epsilon closure for FST2 with cycle detection
        let closure2 = compute_epsilon_closure_for_compose(fst2, s2)?;

        // handle final states (including through epsilon closure)
        for (cs1, w1_eps) in &closure1 {
            for (cs2, w2_eps) in &closure2 {
                if let (Some(w1), Some(w2)) = (fst1.final_weight(*cs1), fst2.final_weight(*cs2)) {
                    let total_weight = w1_eps.times(w1).times(w2_eps).times(w2);
                    if let Some(existing) = result.final_weight(current) {
                        result.set_final(current, existing.plus(&total_weight));
                    } else {
                        result.set_final(current, total_weight);
                    }
                }
            }
        }

        // process arc pairs (including from epsilon closure)
        // Use a map to combine weights for duplicate arcs
        // Track which state pairs in the closure we create arcs from
        #[derive(Clone, PartialEq, Eq, Hash)]
        struct ArcMapKey<FS> {
            ilabel: u32,
            olabel: u32,
            next_state1: StateId,
            next_state2: StateId,
            filter_state: FS,
        }

        let mut arc_map: FxHashMap<ArcMapKey<CF::FilterState>, W> = FxHashMap::default();
        let mut created_arcs_from = FxHashSet::default();
        let mut created_epsilon_from_current = false;

        for (cs1, w1_eps) in &closure1 {
            for (cs2, w2_eps) in &closure2 {
                let mut created_from_this_pair = false;
                for arc1 in fst1.arcs(*cs1) {
                    for arc2 in fst2.arcs(*cs2) {
                        // Match all arcs, including epsilon arcs
                        // The filter will handle matching epsilon arcs correctly
                        if let Some((arc, next_fs)) = filter.filter_arc(&arc1, &arc2, &fs) {
                            created_from_this_pair = true;
                            // Track if we created epsilon arcs from current state
                            if arc.is_epsilon() && (*cs1, *cs2) == (s1, s2) {
                                created_epsilon_from_current = true;
                            }
                            // Accumulate epsilon weights along the path
                            let total_weight = w1_eps.times(&arc.weight).times(w2_eps);
                            let arc_key = ArcMapKey {
                                ilabel: arc.ilabel,
                                olabel: arc.olabel,
                                next_state1: arc1.nextstate,
                                next_state2: arc2.nextstate,
                                filter_state: next_fs.clone(),
                            };

                            // Combine weights for duplicate arcs
                            match arc_map.get_mut(&arc_key) {
                                Some(existing_weight) => {
                                    *existing_weight = existing_weight.plus(&total_weight);
                                }
                                None => {
                                    arc_map.insert(arc_key, total_weight);
                                }
                            }
                        }
                    }
                }
                if created_from_this_pair && (*cs1, *cs2) != (s1, s2) {
                    created_arcs_from.insert((*cs1, *cs2));
                    // Mark as processed since we've created arcs from it via closure
                    processed_via_closure.insert((*cs1, *cs2));
                }
            }
        }

        // Add all arcs from the map
        for (arc_key, weight) in arc_map {
            let next_key = (
                arc_key.next_state1,
                arc_key.next_state2,
                arc_key.filter_state,
            );
            let next_state = match state_map.get(&next_key) {
                Some(&state) => state,
                None => {
                    let state = result.add_state();
                    state_map.insert(next_key.clone(), state);
                    queue.push((next_key.0, next_key.1, next_key.2, state));
                    state
                }
            };

            result.add_arc(
                current,
                Arc::new(arc_key.ilabel, arc_key.olabel, weight, next_state),
            );
        }

        // Mark current state pair as processed via closure
        processed_via_closure.insert((s1, s2));

        // Handle epsilon transitions where only one FST advances
        // Epsilon output from FST1: advance FST1, stay in FST2
        // This handles arcs where arc1.olabel == 0 but arc1.ilabel != 0
        for arc1 in fst1.arcs(s1) {
            if arc1.is_epsilon_output() && !arc1.is_epsilon_input() {
                let next_key = (arc1.nextstate, s2, CF::start());

                // Only process if this state pair hasn't been processed via closure
                // and we haven't created arcs from it via closure
                if !processed_via_closure.contains(&(arc1.nextstate, s2))
                    && !created_arcs_from.contains(&(arc1.nextstate, s2))
                {
                    let next_state = match state_map.get(&next_key) {
                        Some(&state) => state,
                        None => {
                            let state = result.add_state();
                            state_map.insert(next_key.clone(), state);
                            queue.push((arc1.nextstate, s2, CF::start(), state));
                            state
                        }
                    };

                    // Composed arc: input from arc1, output is epsilon (since FST2 didn't advance)
                    let composed_arc = Arc::new(
                        arc1.ilabel,
                        0, // epsilon output
                        arc1.weight.clone(),
                        next_state,
                    );
                    result.add_arc(current, composed_arc);
                }
            }
        }

        // Epsilon input in FST2: stay in FST1, advance FST2
        // This handles arcs where arc2.ilabel == 0 but arc2.olabel != 0
        for arc2 in fst2.arcs(s2) {
            if arc2.is_epsilon_input() && !arc2.is_epsilon_output() {
                let next_key = (s1, arc2.nextstate, CF::start());

                // Only process if this state pair hasn't been processed via closure
                // and we haven't created arcs from it via closure
                if !processed_via_closure.contains(&(s1, arc2.nextstate))
                    && !created_arcs_from.contains(&(s1, arc2.nextstate))
                {
                    let next_state = match state_map.get(&next_key) {
                        Some(&state) => state,
                        None => {
                            let state = result.add_state();
                            state_map.insert(next_key.clone(), state);
                            queue.push((s1, arc2.nextstate, CF::start(), state));
                            state
                        }
                    };

                    // Composed arc: input is epsilon (since FST1 didn't advance), output from arc2
                    let composed_arc = Arc::new(
                        0, // epsilon input
                        arc2.olabel,
                        arc2.weight.clone(),
                        next_state,
                    );
                    result.add_arc(current, composed_arc);
                }
            }
        }

        // Handle full epsilon arcs (both input and output are 0)
        // Full epsilon from FST1: advance FST1, stay in FST2
        // Only handle if we haven't created epsilon arcs from current state via closure
        for arc1 in fst1.arcs(s1) {
            if arc1.is_epsilon() {
                let next_key = (arc1.nextstate, s2, CF::start());

                // Only process if we haven't created epsilon arcs from current state via closure
                // and this state pair hasn't been processed via closure
                // This prevents creating duplicate epsilon arcs
                if !created_epsilon_from_current
                    && !processed_via_closure.contains(&(arc1.nextstate, s2))
                    && !created_arcs_from.contains(&(arc1.nextstate, s2))
                {
                    let next_state = match state_map.get(&next_key) {
                        Some(&state) => state,
                        None => {
                            let state = result.add_state();
                            state_map.insert(next_key.clone(), state);
                            queue.push((arc1.nextstate, s2, CF::start(), state));
                            state
                        }
                    };

                    // Composed arc: full epsilon (both input and output are epsilon)
                    let composed_arc = Arc::epsilon(arc1.weight.clone(), next_state);
                    result.add_arc(current, composed_arc);
                }
            }
        }

        // Full epsilon from FST2: stay in FST1, advance FST2
        // Only handle if we haven't created epsilon arcs from current state via closure
        for arc2 in fst2.arcs(s2) {
            if arc2.is_epsilon() {
                let next_key = (s1, arc2.nextstate, CF::start());

                // Only process if we haven't created epsilon arcs from current state via closure
                // and this state pair hasn't been processed via closure
                // This prevents creating duplicate epsilon arcs
                if !created_epsilon_from_current
                    && !processed_via_closure.contains(&(s1, arc2.nextstate))
                    && !created_arcs_from.contains(&(s1, arc2.nextstate))
                {
                    let next_state = match state_map.get(&next_key) {
                        Some(&state) => state,
                        None => {
                            let state = result.add_state();
                            state_map.insert(next_key.clone(), state);
                            queue.push((s1, arc2.nextstate, CF::start(), state));
                            state
                        }
                    };

                    // Composed arc: full epsilon (both input and output are epsilon)
                    let composed_arc = Arc::epsilon(arc2.weight.clone(), next_state);
                    result.add_arc(current, composed_arc);
                }
            }
        }
    }

    Ok(result)
}

/// Compose two FSTs using the default composition filter
///
/// This is a convenience function that uses [`DefaultComposeFilter`] for standard
/// label-based composition. Use this for most common composition scenarios where
/// the output labels of the first FST should match the input labels of the second FST.
///
/// # Complexity
///
/// - **Time:** O(V₁V₂E₁E₂) worst case, O(VE) average case
///   - V₁, V₂ = number of states in each FST
///   - E₁, E₂ = number of arcs in each FST
/// - **Space:** O(V₁V₂) for reachable state pairs
///
/// # Algorithm
///
/// Delegates to [`compose`] with [`DefaultComposeFilter`], which matches arcs where
/// `arc1.olabel == arc2.ilabel` and combines weights via semiring multiplication (⊗).
///
/// # Examples
///
/// ```rust
/// use arcweight::prelude::*;
///
/// // Build a simple translation pipeline
/// let mut word_to_phone = VectorFst::<TropicalWeight>::new();
/// let s0 = word_to_phone.add_state();
/// let s1 = word_to_phone.add_state();
/// word_to_phone.set_start(s0);
/// word_to_phone.set_final(s1, TropicalWeight::one());
/// word_to_phone.add_arc(s0, Arc::new(1, 2, TropicalWeight::one(), s1));
///
/// let mut phone_to_french = VectorFst::<TropicalWeight>::new();
/// let s0 = phone_to_french.add_state();
/// let s1 = phone_to_french.add_state();
/// phone_to_french.set_start(s0);
/// phone_to_french.set_final(s1, TropicalWeight::one());
/// phone_to_french.add_arc(s0, Arc::new(2, 3, TropicalWeight::one(), s1));
///
/// // Compose: English -> French
/// let translator: VectorFst<TropicalWeight> =
///     compose_default(&word_to_phone, &phone_to_french)?;
/// # Ok::<(), Box<dyn std::error::Error>>(())
/// ```
///
/// # Errors
///
/// Returns [`Error::Algorithm`] if:
/// - Either input FST is invalid or has no start state
/// - Memory allocation fails during computation  
/// - The composition operation encounters incompatible state structures
///
/// # See Also
///
/// - [`compose`] - Full composition with custom filters
/// - [`ComposeFilter`] - Trait for custom composition logic
/// - [`DefaultComposeFilter`] - The filter used by this function
/// - [`determinize`] - Often applied after composition
/// - [`minimize`] - Reduce composed FST size
///
/// [`determinize`]: crate::algorithms::determinize::determinize
/// [`minimize`]: crate::algorithms::minimize::minimize
pub fn compose_default<W, F1, F2, M>(fst1: &F1, fst2: &F2) -> Result<M>
where
    W: Semiring,
    F1: Fst<W>,
    F2: Fst<W>,
    M: MutableFst<W> + Default,
{
    compose(fst1, fst2, DefaultComposeFilter)
}

/// Optimized composition using sorted arc matching
///
/// This version pre-sorts arcs by label for faster matching using the two-pointer
/// technique. For FSTs with many arcs per state, this reduces the matching complexity
/// from O(E₁ × E₂) to O(E₁ + E₂) per state pair.
///
/// # When to Use
///
/// - FSTs with many arcs per state (>10 arcs)
/// - Non-epsilon FSTs (epsilon handling still uses standard matching)
/// - When composition is the bottleneck
///
/// # Complexity
///
/// - **Time:** O(V₁V₂(E₁ log E₁ + E₂ log E₂)) - sorting dominates
/// - **Space:** O(V₁V₂ + max(E₁, E₂)) - state pairs + sorted arc cache
///
/// # Examples
///
/// ```rust
/// use arcweight::prelude::*;
/// use arcweight::algorithms::compose_sorted;
///
/// 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, Arc::new(1, 2, TropicalWeight::new(0.5), s1));
///
/// let mut fst2 = VectorFst::<TropicalWeight>::new();
/// let t0 = fst2.add_state();
/// let t1 = fst2.add_state();
/// fst2.set_start(t0);
/// fst2.set_final(t1, TropicalWeight::one());
/// fst2.add_arc(t0, Arc::new(2, 3, TropicalWeight::new(0.3), t1));
///
/// let composed: VectorFst<TropicalWeight> = compose_sorted(&fst1, &fst2).unwrap();
/// ```
pub fn compose_sorted<W, F1, F2, M>(fst1: &F1, fst2: &F2) -> Result<M>
where
    W: Semiring,
    F1: Fst<W>,
    F2: Fst<W>,
    M: MutableFst<W> + Default,
{
    let start1 = fst1
        .start()
        .ok_or_else(|| Error::Algorithm("First FST has no start state".into()))?;
    let start2 = fst2
        .start()
        .ok_or_else(|| Error::Algorithm("Second FST has no start state".into()))?;

    let mut result = M::default();
    let mut state_map: FxHashMap<(StateId, StateId), StateId> = FxHashMap::default();
    let mut queue = Vec::new();

    // Pre-sort arcs by output label (FST1) and input label (FST2)
    // Cache sorted arcs per state to avoid re-sorting
    let mut sorted_arcs1: FxHashMap<StateId, Vec<Arc<W>>> = FxHashMap::default();
    let mut sorted_arcs2: FxHashMap<StateId, Vec<Arc<W>>> = FxHashMap::default();

    // Create start state
    let start_state = result.add_state();
    result.set_start(start_state);
    state_map.insert((start1, start2), start_state);
    queue.push((start1, start2, start_state));

    while let Some((s1, s2, current)) = queue.pop() {
        // Handle final states
        if let (Some(w1), Some(w2)) = (fst1.final_weight(s1), fst2.final_weight(s2)) {
            let final_weight = w1.times(w2);
            if let Some(existing) = result.final_weight(current) {
                result.set_final(current, existing.plus(&final_weight));
            } else {
                result.set_final(current, final_weight);
            }
        }

        // Get or create sorted arcs for s1 (sorted by olabel)
        let arcs1 = sorted_arcs1.entry(s1).or_insert_with(|| {
            let mut arcs: Vec<Arc<W>> = fst1.arcs(s1).collect();
            arcs.sort_by_key(|a| a.olabel);
            arcs
        });

        // Get or create sorted arcs for s2 (sorted by ilabel)
        let arcs2 = sorted_arcs2.entry(s2).or_insert_with(|| {
            let mut arcs: Vec<Arc<W>> = fst2.arcs(s2).collect();
            arcs.sort_by_key(|a| a.ilabel);
            arcs
        });

        // Separate epsilon arcs (label 0) from non-epsilon arcs
        // Epsilon arcs need special handling
        let eps1: Vec<_> = arcs1.iter().filter(|a| a.olabel == 0).collect();
        let eps2: Vec<_> = arcs2.iter().filter(|a| a.ilabel == 0).collect();
        let non_eps1: Vec<_> = arcs1.iter().filter(|a| a.olabel != 0).collect();
        let non_eps2: Vec<_> = arcs2.iter().filter(|a| a.ilabel != 0).collect();

        // Two-pointer merge for non-epsilon arcs (O(E₁ + E₂))
        let mut i = 0;
        let mut j = 0;
        while i < non_eps1.len() && j < non_eps2.len() {
            let arc1 = non_eps1[i];
            let arc2 = non_eps2[j];

            match arc1.olabel.cmp(&arc2.ilabel) {
                std::cmp::Ordering::Less => i += 1,
                std::cmp::Ordering::Greater => j += 1,
                std::cmp::Ordering::Equal => {
                    // Found matching labels - process all pairs with same labels
                    let label = arc1.olabel;

                    // Find extent of arcs with same label in both lists
                    let mut i_end = i;
                    while i_end < non_eps1.len() && non_eps1[i_end].olabel == label {
                        i_end += 1;
                    }
                    let mut j_end = j;
                    while j_end < non_eps2.len() && non_eps2[j_end].ilabel == label {
                        j_end += 1;
                    }

                    // Match all pairs with same label
                    for a1 in non_eps1.iter().take(i_end).skip(i) {
                        for a2 in non_eps2.iter().take(j_end).skip(j) {
                            let next_key = (a1.nextstate, a2.nextstate);
                            let next_state = match state_map.get(&next_key) {
                                Some(&state) => state,
                                None => {
                                    let state = result.add_state();
                                    state_map.insert(next_key, state);
                                    queue.push((next_key.0, next_key.1, state));
                                    state
                                }
                            };

                            let combined_weight = a1.weight.times(&a2.weight);
                            result.add_arc(
                                current,
                                Arc::new(a1.ilabel, a2.olabel, combined_weight, next_state),
                            );
                        }
                    }

                    i = i_end;
                    j = j_end;
                }
            }
        }

        // Handle epsilon arcs from FST1 (advance FST1 only)
        for arc1 in &eps1 {
            let next_key = (arc1.nextstate, s2);
            let next_state = match state_map.get(&next_key) {
                Some(&state) => state,
                None => {
                    let state = result.add_state();
                    state_map.insert(next_key, state);
                    queue.push((next_key.0, next_key.1, state));
                    state
                }
            };
            result.add_arc(
                current,
                Arc::new(arc1.ilabel, 0, arc1.weight.clone(), next_state),
            );
        }

        // Handle epsilon arcs from FST2 (advance FST2 only)
        for arc2 in &eps2 {
            let next_key = (s1, arc2.nextstate);
            let next_state = match state_map.get(&next_key) {
                Some(&state) => state,
                None => {
                    let state = result.add_state();
                    state_map.insert(next_key, state);
                    queue.push((next_key.0, next_key.1, state));
                    state
                }
            };
            result.add_arc(
                current,
                Arc::new(0, arc2.olabel, arc2.weight.clone(), next_state),
            );
        }
    }

    Ok(result)
}

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

    #[test]
    fn test_basic_composition() {
        // First FST: input -> intermediate
        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, Arc::new(1, 2, TropicalWeight::new(1.0), s1));

        // Second FST: intermediate -> output
        let mut fst2 = VectorFst::<TropicalWeight>::new();
        let t0 = fst2.add_state();
        let t1 = fst2.add_state();
        fst2.set_start(t0);
        fst2.set_final(t1, TropicalWeight::one());
        fst2.add_arc(t0, Arc::new(2, 3, TropicalWeight::new(2.0), t1));

        let composed: VectorFst<TropicalWeight> = compose_default(&fst1, &fst2).unwrap();

        assert!(composed.start().is_some());
        assert!(composed.num_states() > 0);

        // Should have at least one path from input 1 to output 3
        let mut found_path = false;
        for state in composed.states() {
            for arc in composed.arcs(state) {
                if arc.ilabel == 1 && arc.olabel == 3 {
                    found_path = true;
                    // Weight should be combined: 1.0 + 2.0 = 3.0
                    assert_eq!(*arc.weight.value(), 3.0);
                }
            }
        }
        assert!(found_path);
    }

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

        let mut fst2 = VectorFst::<TropicalWeight>::new();
        let t0 = fst2.add_state();
        let t1 = fst2.add_state();
        fst2.set_start(t0);
        fst2.set_final(t1, TropicalWeight::one());
        fst2.add_arc(t0, Arc::new(3, 4, TropicalWeight::new(2.0), t1)); // No match

        let composed: VectorFst<TropicalWeight> = compose_default(&fst1, &fst2).unwrap();

        // Should result in empty FST or FST with no accepting paths
        assert!(composed.start().is_some()); // May have start state

        // Check if any state is final
        let mut has_final = false;
        for state in composed.states() {
            if composed.is_final(state) {
                has_final = true;
                break;
            }
        }
        // Should have no final states reachable from start
        if has_final {
            // If there are final states, they should not be reachable
            assert_eq!(composed.num_arcs_total(), 0);
        }
    }

    #[test]
    fn test_composition_epsilon() {
        // Test case 1: Epsilon output from FST1 matching non-epsilon input in FST2
        let mut fst1 = VectorFst::<TropicalWeight>::new();
        let s0 = fst1.add_state();
        let s1 = fst1.add_state();
        let s2 = fst1.add_state();
        fst1.set_start(s0);
        fst1.set_final(s2, TropicalWeight::one());
        fst1.add_arc(s0, Arc::epsilon(TropicalWeight::new(0.5), s1));
        fst1.add_arc(s1, Arc::new(1, 2, TropicalWeight::new(1.0), s2));

        let mut fst2 = VectorFst::<TropicalWeight>::new();
        let t0 = fst2.add_state();
        let t1 = fst2.add_state();
        fst2.set_start(t0);
        fst2.set_final(t1, TropicalWeight::one());
        fst2.add_arc(t0, Arc::new(2, 3, TropicalWeight::new(2.0), t1));

        let composed: VectorFst<TropicalWeight> = compose_default(&fst1, &fst2).unwrap();

        assert!(composed.start().is_some());
        assert!(composed.num_states() > 0);

        // Should have a path from input 1 to output 3
        let mut found_path = false;
        for state in composed.states() {
            for arc in composed.arcs(state) {
                if arc.ilabel == 1 && arc.olabel == 3 {
                    found_path = true;
                    // Weight should be: 0.5 (epsilon) + 1.0 (arc1) + 2.0 (arc2) = 3.5
                    assert_eq!(*arc.weight.value(), 3.5);
                }
            }
        }
        assert!(found_path, "Should find path from input 1 to output 3");
    }

    #[test]
    fn test_composition_epsilon_input_fst2() {
        // Test case 2: Epsilon input in FST2 matching non-epsilon output from FST1
        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, Arc::new(1, 2, TropicalWeight::new(1.0), s1));

        let mut fst2 = VectorFst::<TropicalWeight>::new();
        let t0 = fst2.add_state();
        let t1 = fst2.add_state();
        let t2 = fst2.add_state();
        fst2.set_start(t0);
        fst2.set_final(t2, TropicalWeight::one());
        fst2.add_arc(t0, Arc::new(0, 0, TropicalWeight::new(0.5), t1)); // epsilon input
        fst2.add_arc(t1, Arc::new(2, 3, TropicalWeight::new(2.0), t2));

        let composed: VectorFst<TropicalWeight> = compose_default(&fst1, &fst2).unwrap();

        assert!(composed.start().is_some());
        assert!(composed.num_states() > 0);

        // Should have a path from input 1 to output 3
        let mut found_path = false;
        for state in composed.states() {
            for arc in composed.arcs(state) {
                if arc.ilabel == 1 && arc.olabel == 3 {
                    found_path = true;
                    // Weight should be: 1.0 (arc1) + 0.5 (epsilon) + 2.0 (arc2) = 3.5
                    assert_eq!(*arc.weight.value(), 3.5);
                }
            }
        }
        assert!(found_path, "Should find path from input 1 to output 3");
    }

    #[test]
    fn test_composition_both_epsilon() {
        // Test case 3: Both FSTs have epsilon transitions that match
        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, Arc::epsilon(TropicalWeight::new(0.5), s1));

        let mut fst2 = VectorFst::<TropicalWeight>::new();
        let t0 = fst2.add_state();
        let t1 = fst2.add_state();
        fst2.set_start(t0);
        fst2.set_final(t1, TropicalWeight::one());
        fst2.add_arc(t0, Arc::epsilon(TropicalWeight::new(0.3), t1));

        let composed: VectorFst<TropicalWeight> = compose_default(&fst1, &fst2).unwrap();

        assert!(composed.start().is_some());
        assert!(composed.num_states() > 0);

        // Should have epsilon path
        let mut found_epsilon = false;
        for state in composed.states() {
            for arc in composed.arcs(state) {
                if arc.is_epsilon() {
                    found_epsilon = true;
                    // Weight should be: 0.5 + 0.3 = 0.8
                    assert_eq!(*arc.weight.value(), 0.8);
                }
            }
        }
        assert!(found_epsilon, "Should find epsilon arc in composition");
    }

    #[test]
    fn test_composition_epsilon_output_advances_fst1() {
        // Test case 4: Epsilon output from FST1 (olabel=0, ilabel!=0) should advance FST1 only
        let mut fst1 = VectorFst::<TropicalWeight>::new();
        let s0 = fst1.add_state();
        let s1 = fst1.add_state();
        let s2 = fst1.add_state();
        fst1.set_start(s0);
        fst1.set_final(s2, TropicalWeight::one());
        // Arc with epsilon output but non-epsilon input
        fst1.add_arc(s0, Arc::new(1, 0, TropicalWeight::new(0.5), s1));
        fst1.add_arc(s1, Arc::new(2, 3, TropicalWeight::new(1.0), s2));

        let mut fst2 = VectorFst::<TropicalWeight>::new();
        let t0 = fst2.add_state();
        let t1 = fst2.add_state();
        fst2.set_start(t0);
        fst2.set_final(t1, TropicalWeight::one());
        fst2.add_arc(t0, Arc::new(3, 4, TropicalWeight::new(2.0), t1));

        let composed: VectorFst<TropicalWeight> = compose_default(&fst1, &fst2).unwrap();

        assert!(composed.start().is_some());

        // Should have path: input 1 -> epsilon output, then input 2 -> output 4
        // The epsilon output from FST1 should allow FST1 to advance while FST2 stays
        let mut found_epsilon_arc = false;
        let mut found_final_arc = false;
        for state in composed.states() {
            for arc in composed.arcs(state) {
                if arc.ilabel == 1 && arc.olabel == 0 {
                    found_epsilon_arc = true;
                    assert_eq!(*arc.weight.value(), 0.5);
                }
                if arc.ilabel == 2 && arc.olabel == 4 {
                    found_final_arc = true;
                    // Weight: 1.0 (arc1) + 2.0 (arc2) = 3.0
                    assert_eq!(*arc.weight.value(), 3.0);
                }
            }
        }
        assert!(
            found_epsilon_arc || found_final_arc,
            "Should handle epsilon output correctly"
        );
    }

    #[test]
    fn test_composition_with_epsilon_cycle_fst1() {
        // Test composition with epsilon cycle in FST1
        let mut fst1 = VectorFst::<TropicalWeight>::new();
        let s0 = fst1.add_state();
        let s1 = fst1.add_state();
        let s2 = fst1.add_state();
        fst1.set_start(s0);
        fst1.set_final(s2, TropicalWeight::one());
        // Epsilon cycle: s0 -> s1 -> s0
        fst1.add_arc(s0, Arc::epsilon(TropicalWeight::new(0.5), s1));
        fst1.add_arc(s1, Arc::epsilon(TropicalWeight::new(0.3), s0));
        // Exit from cycle
        fst1.add_arc(s0, Arc::new(1, 2, TropicalWeight::new(1.0), s2));

        let mut fst2 = VectorFst::<TropicalWeight>::new();
        let t0 = fst2.add_state();
        let t1 = fst2.add_state();
        fst2.set_start(t0);
        fst2.set_final(t1, TropicalWeight::one());
        fst2.add_arc(t0, Arc::new(2, 3, TropicalWeight::new(2.0), t1));

        let composed: VectorFst<TropicalWeight> = compose_default(&fst1, &fst2).unwrap();

        assert!(composed.start().is_some());
        // Should have path from input 1 to output 3
        let mut found_path = false;
        for state in composed.states() {
            for arc in composed.arcs(state) {
                if arc.ilabel == 1 && arc.olabel == 3 {
                    found_path = true;
                    break;
                }
            }
            if found_path {
                break;
            }
        }
        assert!(found_path, "Should find path from input 1 to output 3");
    }

    #[test]
    fn test_composition_with_epsilon_cycle_fst2() {
        // Test composition with epsilon cycle in FST2
        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, Arc::new(1, 2, TropicalWeight::new(1.0), s1));

        let mut fst2 = VectorFst::<TropicalWeight>::new();
        let t0 = fst2.add_state();
        let t1 = fst2.add_state();
        let t2 = fst2.add_state();
        fst2.set_start(t0);
        fst2.set_final(t2, TropicalWeight::one());
        // Epsilon cycle: t0 -> t1 -> t0
        fst2.add_arc(t0, Arc::epsilon(TropicalWeight::new(0.5), t1));
        fst2.add_arc(t1, Arc::epsilon(TropicalWeight::new(0.3), t0));
        // Exit from cycle
        fst2.add_arc(t0, Arc::new(2, 3, TropicalWeight::new(2.0), t2));

        let composed: VectorFst<TropicalWeight> = compose_default(&fst1, &fst2).unwrap();

        assert!(composed.start().is_some());
        // Should have path from input 1 to output 3
        let mut found_path = false;
        for state in composed.states() {
            for arc in composed.arcs(state) {
                if arc.ilabel == 1 && arc.olabel == 3 {
                    found_path = true;
                    break;
                }
            }
            if found_path {
                break;
            }
        }
        assert!(found_path, "Should find path from input 1 to output 3");
    }

    #[test]
    fn test_composition_with_epsilon_cycles_both() {
        // Test composition with epsilon cycles in both FSTs
        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());
        // Self-loop on s0
        fst1.add_arc(s0, Arc::epsilon(TropicalWeight::new(0.2), s0));
        fst1.add_arc(s0, Arc::new(1, 2, TropicalWeight::new(1.0), s1));

        let mut fst2 = VectorFst::<TropicalWeight>::new();
        let t0 = fst2.add_state();
        let t1 = fst2.add_state();
        fst2.set_start(t0);
        fst2.set_final(t1, TropicalWeight::one());
        // Self-loop on t0
        fst2.add_arc(t0, Arc::epsilon(TropicalWeight::new(0.3), t0));
        fst2.add_arc(t0, Arc::new(2, 3, TropicalWeight::new(2.0), t1));

        let composed: VectorFst<TropicalWeight> = compose_default(&fst1, &fst2).unwrap();

        assert!(composed.start().is_some());
        // Should terminate and produce result
        assert!(composed.num_states() > 0);
    }

    #[test]
    fn test_composition_epsilon_cycle_termination() {
        // Test that composition terminates even with epsilon cycles
        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());
        // Large cycle that could cause infinite loop
        fst1.add_arc(s0, Arc::epsilon(TropicalWeight::new(0.1), s0));
        fst1.add_arc(s0, Arc::new(1, 2, TropicalWeight::new(1.0), s1));

        let mut fst2 = VectorFst::<TropicalWeight>::new();
        let t0 = fst2.add_state();
        let t1 = fst2.add_state();
        fst2.set_start(t0);
        fst2.set_final(t1, TropicalWeight::one());
        fst2.add_arc(t0, Arc::new(2, 3, TropicalWeight::new(1.0), t1));

        // This should terminate quickly
        let composed: VectorFst<TropicalWeight> = compose_default(&fst1, &fst2).unwrap();
        assert!(composed.start().is_some());
    }

    #[test]
    fn test_compose_sorted_basic() {
        // Test compose_sorted produces same results as compose_default for non-epsilon FSTs
        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());
        // Multiple arcs with different labels
        fst1.add_arc(s0, Arc::new(1, 2, TropicalWeight::new(1.0), s1));
        fst1.add_arc(s0, Arc::new(2, 3, TropicalWeight::new(1.5), s1));
        fst1.add_arc(s0, Arc::new(3, 4, TropicalWeight::new(2.0), s1));

        let mut fst2 = VectorFst::<TropicalWeight>::new();
        let t0 = fst2.add_state();
        let t1 = fst2.add_state();
        fst2.set_start(t0);
        fst2.set_final(t1, TropicalWeight::one());
        // Multiple arcs matching some of FST1
        fst2.add_arc(t0, Arc::new(2, 5, TropicalWeight::new(0.5), t1));
        fst2.add_arc(t0, Arc::new(4, 6, TropicalWeight::new(0.3), t1));

        let composed: VectorFst<TropicalWeight> = compose_sorted(&fst1, &fst2).unwrap();

        // Should have paths 1->5 (via 2) and 3->6 (via 4)
        let mut found_path1 = false;
        let mut found_path2 = false;
        for state in composed.states() {
            for arc in composed.arcs(state) {
                if arc.ilabel == 1 && arc.olabel == 5 {
                    found_path1 = true;
                    // Weight: 1.0 + 0.5 = 1.5
                    assert_eq!(*arc.weight.value(), 1.5);
                }
                if arc.ilabel == 3 && arc.olabel == 6 {
                    found_path2 = true;
                    // Weight: 2.0 + 0.3 = 2.3
                    assert_eq!(*arc.weight.value(), 2.3);
                }
            }
        }
        assert!(found_path1, "Should find path 1->5");
        assert!(found_path2, "Should find path 3->6");
    }

    #[test]
    fn test_compose_sorted_many_arcs() {
        // Test with many arcs per state to verify sorting optimization
        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());

        let mut fst2 = VectorFst::<TropicalWeight>::new();
        let t0 = fst2.add_state();
        let t1 = fst2.add_state();
        fst2.set_start(t0);
        fst2.set_final(t1, TropicalWeight::one());

        // Add many arcs with various labels
        for i in 1..=50 {
            fst1.add_arc(
                s0,
                Arc::new(i, i * 2, TropicalWeight::new(i as f32 * 0.1), s1),
            );
        }

        for i in 1..=50 {
            fst2.add_arc(
                t0,
                Arc::new(i * 2, i * 3, TropicalWeight::new(i as f32 * 0.05), t1),
            );
        }

        let composed: VectorFst<TropicalWeight> = compose_sorted(&fst1, &fst2).unwrap();

        // Should have 50 matching paths
        let mut arc_count = 0;
        for state in composed.states() {
            arc_count += composed.num_arcs(state);
        }
        assert_eq!(arc_count, 50, "Should have 50 composed arcs");
    }

    #[test]
    fn test_compose_sorted_with_epsilon() {
        // Test that compose_sorted handles epsilon arcs correctly
        let mut fst1 = VectorFst::<TropicalWeight>::new();
        let s0 = fst1.add_state();
        let s1 = fst1.add_state();
        let s2 = fst1.add_state();
        fst1.set_start(s0);
        fst1.set_final(s2, TropicalWeight::one());
        fst1.add_arc(s0, Arc::epsilon(TropicalWeight::new(0.5), s1)); // epsilon
        fst1.add_arc(s1, Arc::new(1, 2, TropicalWeight::new(1.0), s2));

        let mut fst2 = VectorFst::<TropicalWeight>::new();
        let t0 = fst2.add_state();
        let t1 = fst2.add_state();
        fst2.set_start(t0);
        fst2.set_final(t1, TropicalWeight::one());
        fst2.add_arc(t0, Arc::new(2, 3, TropicalWeight::new(2.0), t1));

        let composed: VectorFst<TropicalWeight> = compose_sorted(&fst1, &fst2).unwrap();

        assert!(composed.start().is_some());
        assert!(composed.num_states() > 0);
    }
}