arcweight 0.3.0

//! Lookahead composition filters for optimized FST composition.
//!
//! Lookahead filters examine future paths during composition to prune non-productive
//! state pairs early, significantly reducing the composed FST's state space. This is
//! particularly effective when the input FSTs have many arcs that cannot lead to
//! successful paths.
//!
//! # Overview
//!
//! Standard composition explores all matching arc pairs, even when they lead to
//! dead ends. Lookahead filters precompute reachability information to skip
//! non-productive compositions:
//!
//! ```text
//! Standard:   Creates state pair (q₁, q₂) → discovers dead end later
//! Lookahead:  Checks future paths → skips (q₁, q₂) immediately
//! ```
//!
//! # Filter Types
//!
//! | Filter | Description | Best For |
//! |--------|-------------|----------|
//! | [`LabelLookaheadFilter`] | Match by label compatibility | Simple compositions |
//! | [`LabelPairLookaheadFilter`] | Examine future (input, output) pairs | Complex pipelines |
//! | [`MatcherLookaheadFilter`] | Custom matching function | Special requirements |
//!
//! # Complexity
//!
//! Lookahead adds preprocessing cost but typically provides substantial savings:
//!
//! - **Preprocessing:** $`O(V + E)`$ per FST for reachability computation
//! - **Per-state overhead:** $`O(1)`$ lookup per state pair
//! - **Space:** $`O(|\Sigma|)`$ for label sets
//!
//! # When to Use
//!
//! Lookahead composition is most effective when:
//! - FSTs have large label alphabets with low overlap
//! - Standard composition produces many unreachable states
//! - Memory is constrained and state count must be minimized
//!
//! # Example
//!
//! ```rust
//! use arcweight::prelude::*;
//! use arcweight::algorithms::{compose_with_lookahead, LabelPairLookaheadFilter};
//!
//! 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::one(), 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::one(), t1));
//!
//! // Use lookahead to optimize composition
//! let filter = LabelPairLookaheadFilter::with_max_depth(2);
//! let composed: VectorFst<TropicalWeight> = compose_with_lookahead(&fst1, &fst2, filter)?;
//! # Ok::<(), arcweight::Error>(())
//! ```
//!
//! # References
//!
//! \[1\] Allauzen, C. and Mohri, M. 2009. N-way composition of weighted finite-state
//!     transducers. *International Journal of Foundations of Computer Science*
//!     20, 4 (August 2009), 613-627. <https://doi.org/10.1142/S0129054109006784>
//!
//! \[2\] Mohri, M., Pereira, F., and Riley, M. 2000. The design principles of a
//!     weighted finite-state transducer library. *Theoretical Computer Science*
//!     231, 1 (January 2000), 17-32. <https://doi.org/10.1016/S0304-3975(99)00014-6>

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

/// Trait for lookahead composition filters that prune non-productive state pairs.
///
/// Lookahead filters extend basic composition by examining future paths before
/// creating composed states. This enables early pruning of state pairs that
/// cannot lead to accepting paths.
///
/// # Required Methods
///
/// - [`filter_arc`](Self::filter_arc) - Determines if two arcs can be composed
///
/// # Optional Methods
///
/// - [`lookahead_path`](Self::lookahead_path) - Checks if destination states have
///   productive future paths (default: always returns `true`)
///
/// # Implementing Custom Filters
///
/// ```rust
/// use arcweight::prelude::*;
/// use arcweight::algorithms::LookaheadComposeFilter;
///
/// struct WeightThresholdFilter {
///     threshold: f32,
/// }
///
/// impl<W: Semiring> LookaheadComposeFilter<W> for WeightThresholdFilter {
///     fn filter_arc(
///         &self,
///         _fst1_state: StateId,
///         _fst2_state: StateId,
///         arc1: &Arc<W>,
///         arc2: &Arc<W>,
///     ) -> bool {
///         // Match labels
///         arc1.olabel == arc2.ilabel
///     }
/// }
/// ```
///
/// # See Also
///
/// - [`LabelLookaheadFilter`] - Basic label matching
/// - [`LabelPairLookaheadFilter`] - Future path examination
/// - [`compose_with_lookahead`] - Main composition entry point
pub trait LookaheadComposeFilter<W: Semiring> {
    /// Determines if two arcs can be composed based on label compatibility.
    ///
    /// # Arguments
    ///
    /// * `fst1_state` - Source state in FST1
    /// * `fst2_state` - Source state in FST2
    /// * `arc1` - Arc from FST1
    /// * `arc2` - Arc from FST2
    ///
    /// # Returns
    ///
    /// `true` if the arcs should be composed, `false` to skip.
    fn filter_arc(
        &self,
        fst1_state: StateId,
        fst2_state: StateId,
        arc1: &Arc<W>,
        arc2: &Arc<W>,
    ) -> bool;

    /// Performs lookahead to check if destination states have productive paths.
    ///
    /// This method is called after `filter_arc` returns `true` to determine
    /// whether to actually create the composed state. Implementations can
    /// examine future arcs to detect dead ends early.
    ///
    /// # Arguments
    ///
    /// * `fst1` - First FST
    /// * `fst2` - Second FST
    /// * `next1` - Destination state in FST1
    /// * `next2` - Destination state in FST2
    ///
    /// # Returns
    ///
    /// `true` if the state pair should be created, `false` to prune.
    ///
    /// # Default Implementation
    ///
    /// Returns `true` (no pruning), suitable for filters that only check
    /// immediate arc compatibility.
    fn lookahead_path<F1: Fst<W>, F2: Fst<W>>(
        &self,
        fst1: &F1,
        fst2: &F2,
        next1: StateId,
        next2: StateId,
    ) -> bool {
        let _ = (fst1, fst2, next1, next2);
        true // Default: no lookahead filtering
    }
}

/// Simple label-matching lookahead filter.
///
/// Matches arcs when the output label of the first arc equals the input label
/// of the second arc, with standard epsilon handling.
///
/// # Matching Rules
///
/// | `arc1.olabel` | `arc2.ilabel` | Match? |
/// |---------------|---------------|--------|
/// | `a` | `a` | Yes |
/// | `a` | `b` | No |
/// | `0` (epsilon) | any | Yes |
/// | any | `0` (epsilon) | Yes |
///
/// # Example
///
/// ```rust
/// use arcweight::prelude::*;
/// use arcweight::algorithms::{compose_with_lookahead, LabelLookaheadFilter};
///
/// let mut fst1 = VectorFst::<TropicalWeight>::new();
/// let s = fst1.add_state();
/// fst1.set_start(s);
/// fst1.set_final(s, TropicalWeight::one());
///
/// let mut fst2 = VectorFst::<TropicalWeight>::new();
/// let t = fst2.add_state();
/// fst2.set_start(t);
/// fst2.set_final(t, TropicalWeight::one());
///
/// let filter = LabelLookaheadFilter;
/// let result: VectorFst<TropicalWeight> = compose_with_lookahead(&fst1, &fst2, filter)?;
/// # Ok::<(), arcweight::Error>(())
/// ```
#[derive(Debug, Clone, Copy, Default)]
pub struct LabelLookaheadFilter;

/// Lookahead filter with future path examination.
///
/// Extends basic label matching by examining future (input, output) label pairs
/// along paths from destination states. Only creates composed states when
/// productive paths exist forward, significantly reducing the composed FST's
/// state space.
///
/// # Algorithm
///
/// For each potential composed state pair $`(q_1, q_2)`$:
///
/// 1. Check immediate label compatibility (same as [`LabelLookaheadFilter`])
/// 2. Collect output labels reachable from $`q_1`$
/// 3. Collect input labels reachable from $`q_2`$
/// 4. Verify non-empty intersection (productive paths exist)
/// 5. Optionally recurse to `max_depth` for deeper analysis
///
/// # Complexity
///
/// - **Per-state check:** $`O(E)`$ for collecting reachable labels
/// - **With recursion:** $`O(E^d)`$ where $`d`$ is `max_depth`
///
/// # Configuration
///
/// - `max_depth` - Maximum recursion depth (default: 2)
///   - Higher values catch more dead ends but increase overhead
///   - Recommended: 1-3 for most applications
///
/// # Example
///
/// ```rust
/// use arcweight::prelude::*;
/// use arcweight::algorithms::{compose_with_lookahead, LabelPairLookaheadFilter};
///
/// 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::one(), 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::one(), t1));
///
/// // Deeper lookahead for complex FSTs
/// let filter = LabelPairLookaheadFilter::with_max_depth(3);
/// let result: VectorFst<TropicalWeight> = compose_with_lookahead(&fst1, &fst2, filter)?;
/// # Ok::<(), arcweight::Error>(())
/// ```
#[derive(Debug, Clone)]
pub struct LabelPairLookaheadFilter {
    /// Maximum depth to examine future paths (default: 2).
    pub max_depth: usize,
}

impl LabelPairLookaheadFilter {
    /// Create a new label pair lookahead filter
    pub fn new() -> Self {
        Self { max_depth: 2 }
    }

    /// Create with custom maximum lookahead depth
    pub fn with_max_depth(max_depth: usize) -> Self {
        Self { max_depth }
    }
}

impl Default for LabelPairLookaheadFilter {
    fn default() -> Self {
        Self::new()
    }
}

impl<W: Semiring> LookaheadComposeFilter<W> for LabelLookaheadFilter {
    fn filter_arc(
        &self,
        _fst1_state: StateId,
        _fst2_state: StateId,
        arc1: &Arc<W>,
        arc2: &Arc<W>,
    ) -> bool {
        // Match output label of arc1 with input label of arc2
        // For composition T₁ ∘ T₂, we need T₁'s output to match T₂'s input
        // Epsilon handling:
        // - Epsilon on FST1 output (arc1.olabel == 0) matches any FST2 input
        // - Any FST1 output matches epsilon on FST2 input (arc2.ilabel == 0)
        // - Epsilon on both sides (both 0) matches
        use crate::fst::NO_LABEL;
        arc1.olabel == arc2.ilabel
            || arc1.olabel == NO_LABEL  // Epsilon on FST1 output matches any FST2 input
            || arc2.ilabel == NO_LABEL // Any FST1 output matches epsilon on FST2 input
    }
}

impl LabelPairLookaheadFilter {
    /// Check if there are compatible future paths from the given states
    ///
    /// This implements label pair lookahead by examining future (input, output)
    /// label pairs along paths. Returns true if there are compatible paths
    /// that could lead to productive composition.
    fn has_compatible_future_paths<W: Semiring, F1: Fst<W>, F2: Fst<W>>(
        &self,
        fst1: &F1,
        fst2: &F2,
        state1: StateId,
        state2: StateId,
        depth: usize,
    ) -> bool {
        // Limit lookahead depth to avoid excessive computation
        if depth >= self.max_depth {
            return true; // Assume productive at max depth
        }

        // Collect output labels from fst1 and input labels from fst2
        let mut fst1_outputs = std::collections::HashSet::new();
        let mut fst2_inputs = std::collections::HashSet::new();
        let mut fst1_has_epsilon_output = false;
        let mut fst2_has_epsilon_input = false;

        // Collect output labels from fst1
        for arc in fst1.arcs(state1) {
            if arc.is_epsilon() {
                fst1_has_epsilon_output = true;
            } else {
                fst1_outputs.insert(arc.olabel);
            }
        }

        // Collect input labels from fst2
        for arc in fst2.arcs(state2) {
            if arc.is_epsilon() {
                fst2_has_epsilon_input = true;
            } else {
                fst2_inputs.insert(arc.ilabel);
            }
        }

        // Check for compatible label pairs (fst1 output matches fst2 input)
        // Epsilon handling: epsilon on either side makes composition possible
        let has_compatible = fst1_outputs.intersection(&fst2_inputs).next().is_some()
            || fst1_has_epsilon_output  // Epsilon on FST1 output matches any FST2 input
            || fst2_has_epsilon_input; // Any FST1 output matches epsilon on FST2 input

        if !has_compatible {
            return false;
        }

        // If we have compatible labels, check if paths lead to productive states
        // by examining one level deeper
        if depth + 1 < self.max_depth {
            for arc1 in fst1.arcs(state1) {
                for arc2 in fst2.arcs(state2) {
                    // Check if arcs are compatible (including epsilon handling)
                    let is_compatible = if arc1.is_epsilon() {
                        // Epsilon on FST1 output matches any FST2 input
                        true
                    } else if arc2.is_epsilon() {
                        // Any FST1 output matches epsilon on FST2 input
                        true
                    } else {
                        // Normal matching: fst1 output must equal fst2 input
                        arc1.olabel == arc2.ilabel
                    };

                    if is_compatible {
                        // Check if destination states have productive paths
                        if fst1.final_weight(arc1.nextstate).is_some()
                            || fst2.final_weight(arc2.nextstate).is_some()
                        {
                            return true;
                        }

                        // Recursively check deeper
                        if self.has_compatible_future_paths(
                            fst1,
                            fst2,
                            arc1.nextstate,
                            arc2.nextstate,
                            depth + 1,
                        ) {
                            return true;
                        }
                    }
                }
            }
        }

        // Found compatible labels
        true
    }
}

impl<W: Semiring> LookaheadComposeFilter<W> for LabelPairLookaheadFilter {
    fn filter_arc(
        &self,
        _fst1_state: StateId,
        _fst2_state: StateId,
        arc1: &Arc<W>,
        arc2: &Arc<W>,
    ) -> bool {
        // Match output label of arc1 with input label of arc2
        // Epsilon handling:
        // - Epsilon on FST1 output (arc1.olabel == 0) matches any FST2 input
        // - Any FST1 output matches epsilon on FST2 input (arc2.ilabel == 0)
        // - Epsilon on both sides (both 0) matches
        use crate::fst::NO_LABEL;
        arc1.olabel == arc2.ilabel
            || arc1.olabel == NO_LABEL  // Epsilon on FST1 output matches any FST2 input
            || arc2.ilabel == NO_LABEL // Any FST1 output matches epsilon on FST2 input
    }

    fn lookahead_path<F1: Fst<W>, F2: Fst<W>>(
        &self,
        fst1: &F1,
        fst2: &F2,
        next1: StateId,
        next2: StateId,
    ) -> bool {
        // Check if either state is final (productive path)
        if fst1.final_weight(next1).is_some() || fst2.final_weight(next2).is_some() {
            return true;
        }

        // Examine future label pairs along paths
        self.has_compatible_future_paths(fst1, fst2, next1, next2, 0)
    }
}

/// Custom matcher lookahead filter using a user-provided function.
///
/// Enables custom arc matching logic beyond simple label comparison.
/// Useful for specialized composition requirements such as:
///
/// - Weight-based filtering
/// - Multi-tape transducer composition
/// - Constraint propagation
///
/// # Type Parameters
///
/// * `W` - Weight type (must implement [`Semiring`])
/// * `F` - Matcher function type: `Fn(&Arc<W>, &Arc<W>) -> bool`
///
/// # Example
///
/// ```rust
/// use arcweight::prelude::*;
/// use arcweight::algorithms::{compose_with_lookahead, MatcherLookaheadFilter};
///
/// // Only compose arcs with matching weights
/// let filter = MatcherLookaheadFilter::new(
///     |arc1: &Arc<TropicalWeight>, arc2: &Arc<TropicalWeight>| {
///         arc1.olabel == arc2.ilabel && arc1.weight == arc2.weight
///     }
/// );
///
/// let mut fst1 = VectorFst::<TropicalWeight>::new();
/// let s = fst1.add_state();
/// fst1.set_start(s);
/// fst1.set_final(s, TropicalWeight::one());
///
/// let mut fst2 = VectorFst::<TropicalWeight>::new();
/// let t = fst2.add_state();
/// fst2.set_start(t);
/// fst2.set_final(t, TropicalWeight::one());
///
/// let result: VectorFst<TropicalWeight> = compose_with_lookahead(&fst1, &fst2, filter)?;
/// # Ok::<(), arcweight::Error>(())
/// ```
#[derive(Debug)]
pub struct MatcherLookaheadFilter<W: Semiring, F: Fn(&Arc<W>, &Arc<W>) -> bool> {
    matcher: F,
    _phantom: std::marker::PhantomData<W>,
}

impl<W: Semiring, F: Fn(&Arc<W>, &Arc<W>) -> bool> MatcherLookaheadFilter<W, F> {
    /// Create a new matcher lookahead filter with a custom matcher function
    pub fn new(matcher: F) -> Self {
        Self {
            matcher,
            _phantom: std::marker::PhantomData,
        }
    }
}

impl<W: Semiring, F: Fn(&Arc<W>, &Arc<W>) -> bool> LookaheadComposeFilter<W>
    for MatcherLookaheadFilter<W, F>
{
    fn filter_arc(
        &self,
        _fst1_state: StateId,
        _fst2_state: StateId,
        arc1: &Arc<W>,
        arc2: &Arc<W>,
    ) -> bool {
        (self.matcher)(arc1, arc2)
    }
    // Default lookahead_path implementation (returns true) is used
}

/// Composes two FSTs using a lookahead filter for optimization.
///
/// Performs FST composition while applying a lookahead filter to prune
/// non-productive state pairs during construction. This can significantly
/// reduce the composed FST's size when the input FSTs have many incompatible
/// label combinations.
///
/// # Type Parameters
///
/// * `W` - Weight type implementing [`Semiring`]
/// * `F1` - First FST type
/// * `F2` - Second FST type
/// * `M` - Output FST type (mutable)
/// * `Filter` - Lookahead filter implementing [`LookaheadComposeFilter`]
///
/// # Arguments
///
/// * `fst1` - First FST (provides input labels)
/// * `fst2` - Second FST (provides output labels)
/// * `filter` - Lookahead filter for arc matching and path pruning
///
/// # Returns
///
/// The composed FST $`T_1 \circ T_2`$, potentially smaller than standard
/// composition due to lookahead pruning.
///
/// # Complexity
///
/// - **Time:** $`O(V_1 V_2 (E_1 + E_2) + L)`$ where $`L`$ is lookahead cost
/// - **Space:** $`O(V_1 V_2)`$ for state pair tracking
///
/// # Algorithm
///
/// 1. Initialize with start state pair $`(q_{01}, q_{02})`$
/// 2. For each state pair in queue:
///    a. For each arc pair, apply `filter.filter_arc()`
///    b. For matching arcs, apply `filter.lookahead_path()`
///    c. Create composed arc only if both checks pass
/// 3. Handle epsilon transitions with special lookahead
///
/// # Examples
///
/// ## Basic Usage
///
/// ```rust
/// use arcweight::prelude::*;
/// use arcweight::algorithms::{compose_with_lookahead, LabelLookaheadFilter};
///
/// 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, 1, TropicalWeight::one(), s1));
///
/// 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(1, 2, TropicalWeight::one(), s1));
///
/// let filter = LabelLookaheadFilter;
/// let result: VectorFst<TropicalWeight> = compose_with_lookahead(&fst1, &fst2, filter)?;
/// assert!(result.num_states() > 0);
/// # Ok::<(), arcweight::Error>(())
/// ```
///
/// ## With Path Lookahead
///
/// ```rust
/// use arcweight::prelude::*;
/// use arcweight::algorithms::{compose_with_lookahead, LabelPairLookaheadFilter};
///
/// # let mut fst1 = VectorFst::<TropicalWeight>::new();
/// # let s = fst1.add_state();
/// # fst1.set_start(s);
/// # fst1.set_final(s, TropicalWeight::one());
/// # let mut fst2 = fst1.clone();
/// // Use deeper lookahead for better pruning
/// let filter = LabelPairLookaheadFilter::with_max_depth(3);
/// let result: VectorFst<TropicalWeight> = compose_with_lookahead(&fst1, &fst2, filter)?;
/// # Ok::<(), arcweight::Error>(())
/// ```
///
/// # Errors
///
/// Returns [`Error::Algorithm`](crate::Error::Algorithm) if:
/// - Either FST has no start state
/// - Memory allocation fails during composition
///
/// # References
///
/// \[1\] Allauzen, C. and Mohri, M. 2009. N-way composition of weighted finite-state
///     transducers. *International Journal of Foundations of Computer Science*
///     20, 4, 613-627. <https://doi.org/10.1142/S0129054109006784>
///
/// # See Also
///
/// - [`compose`](crate::algorithms::compose()) - Standard composition without lookahead
/// - [`LookaheadComposeFilter`] - Trait for custom filters
pub fn compose_with_lookahead<W, F1, F2, M, Filter>(
    fst1: &F1,
    fst2: &F2,
    filter: Filter,
) -> Result<M>
where
    W: Semiring + Clone,
    F1: Fst<W>,
    F2: Fst<W>,
    M: crate::fst::MutableFst<W> + Default,
    Filter: LookaheadComposeFilter<W>,
{
    use std::collections::HashMap;

    let start1 = fst1
        .start()
        .ok_or_else(|| crate::Error::Algorithm("First FST has no start state".into()))?;
    let start2 = fst2
        .start()
        .ok_or_else(|| crate::Error::Algorithm("Second FST has no start state".into()))?;

    let mut result = M::default();
    let mut state_map = HashMap::new();
    let mut queue = Vec::new();

    // 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));

    // process states
    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)) {
            result.set_final(current, w1.times(w2));
        }

        use crate::fst::NO_LABEL;

        // Process arc pairs with lookahead (including epsilon handling)
        for arc1 in fst1.arcs(s1) {
            for arc2 in fst2.arcs(s2) {
                // Use lookahead filter with actual state IDs
                // The filter can examine arcs and states to determine compatibility
                if filter.filter_arc(s1, s2, &arc1, &arc2) {
                    let next1 = arc1.nextstate;
                    let next2 = arc2.nextstate;

                    // Perform comprehensive lookahead to check if this path is productive
                    if !filter.lookahead_path(fst1, fst2, next1, next2) {
                        continue; // Skip this arc - no productive path forward
                    }

                    let next_key = (next1, next2);

                    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((next1, next2, state));
                            state
                        }
                    };

                    // Create composed arc with correct label matching and weight combination
                    // For composition T₁ ∘ T₂: input from T₁, output from T₂, weight = w₁ ⊗ w₂
                    // Epsilon handling:
                    // - If arc1.olabel is epsilon: composed arc has ilabel from arc1, olabel from arc2
                    // - If arc2.ilabel is epsilon: composed arc has ilabel from arc1, olabel from arc2
                    // - If both are epsilon: composed arc is epsilon
                    let composed_ilabel = if arc1.olabel == NO_LABEL && arc2.ilabel == NO_LABEL {
                        NO_LABEL // Both epsilon -> epsilon
                    } else {
                        arc1.ilabel // Input from FST1
                    };
                    let composed_olabel = if arc1.olabel == NO_LABEL && arc2.ilabel == NO_LABEL {
                        NO_LABEL // Both epsilon -> epsilon
                    } else {
                        arc2.olabel // Output from FST2
                    };

                    let composed_arc = crate::arc::Arc::new(
                        composed_ilabel,
                        composed_olabel,
                        arc1.weight.times(&arc2.weight),
                        next_state,
                    );
                    result.add_arc(current, composed_arc);
                }
            }
        }

        // Handle epsilon transitions where only one FST advances
        // Epsilon on FST1 output: advance FST1, stay in FST2
        for arc1 in fst1.arcs(s1) {
            if arc1.is_epsilon() {
                let next1 = arc1.nextstate;
                let next_key = (next1, s2); // Advance FST1, stay in FST2

                // Check if this path is productive
                if !filter.lookahead_path(fst1, fst2, next1, s2) {
                    continue;
                }

                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((next1, s2, state));
                        state
                    }
                };

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

        // Epsilon on FST2 input: stay in FST1, advance FST2
        for arc2 in fst2.arcs(s2) {
            if arc2.is_epsilon() {
                let next2 = arc2.nextstate;
                let next_key = (s1, next2); // Stay in FST1, advance FST2

                // Check if this path is productive
                if !filter.lookahead_path(fst1, fst2, s1, next2) {
                    continue;
                }

                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((s1, next2, state));
                        state
                    }
                };

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

    Ok(result)
}

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

    #[test]
    fn test_label_lookahead_filter() {
        let filter = LabelLookaheadFilter;
        // Test matching: arc1.olabel (1) == arc2.ilabel (1)
        let arc1 = Arc::new(1, 1, TropicalWeight::one(), 0);
        let arc2 = Arc::new(1, 1, TropicalWeight::one(), 0);
        assert!(filter.filter_arc(0, 0, &arc1, &arc2));

        // Test non-matching: arc1.olabel (1) != arc2.ilabel (3)
        let arc3 = Arc::new(3, 2, TropicalWeight::one(), 0);
        assert!(!filter.filter_arc(0, 0, &arc1, &arc3));

        // Test matching: arc1.olabel (2) == arc2.ilabel (2)
        let arc4 = Arc::new(1, 2, TropicalWeight::one(), 0);
        let arc5 = Arc::new(2, 3, TropicalWeight::one(), 0);
        assert!(filter.filter_arc(0, 0, &arc4, &arc5));

        // Test epsilon handling: epsilon on FST1 output matches any FST2 input
        let epsilon_arc1 = Arc::epsilon(TropicalWeight::one(), 0);
        let arc6 = Arc::new(1, 2, TropicalWeight::one(), 0);
        assert!(filter.filter_arc(0, 0, &epsilon_arc1, &arc6));

        // Test epsilon handling: any FST1 output matches epsilon on FST2 input
        let arc7 = Arc::new(1, 2, TropicalWeight::one(), 0);
        let epsilon_arc2 = Arc::epsilon(TropicalWeight::one(), 0);
        assert!(filter.filter_arc(0, 0, &arc7, &epsilon_arc2));

        // Test epsilon handling: both epsilon match
        assert!(filter.filter_arc(0, 0, &epsilon_arc1, &epsilon_arc2));
    }

    #[test]
    fn test_matcher_lookahead_filter() {
        let filter: MatcherLookaheadFilter<TropicalWeight, _> = MatcherLookaheadFilter::new(
            |arc1: &Arc<TropicalWeight>, arc2: &Arc<TropicalWeight>| {
                arc1.olabel == arc2.ilabel && *arc1.weight.value() == *arc2.weight.value()
            },
        );

        let arc1 = Arc::new(1, 1, TropicalWeight::new(1.0), 0);
        let arc2 = Arc::new(1, 1, TropicalWeight::new(1.0), 0);
        assert!(filter.filter_arc(0, 0, &arc1, &arc2));

        let arc3 = Arc::new(1, 1, TropicalWeight::new(2.0), 0);
        assert!(!filter.filter_arc(0, 0, &arc1, &arc3));
    }

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

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

        let filter = LabelLookaheadFilter;
        let result: VectorFst<TropicalWeight> =
            compose_with_lookahead(&fst1, &fst2, filter).unwrap();
        assert!(result.num_states() > 0);
    }

    #[test]
    fn test_label_pair_lookahead_filter() {
        let filter = LabelPairLookaheadFilter::new();

        // Test basic label matching
        let arc1 = Arc::new(1, 1, TropicalWeight::one(), 0);
        let arc2 = Arc::new(1, 1, TropicalWeight::one(), 0);
        assert!(filter.filter_arc(0, 0, &arc1, &arc2));

        let arc3 = Arc::new(3, 2, TropicalWeight::one(), 0);
        assert!(!filter.filter_arc(0, 0, &arc1, &arc3));

        // Test epsilon handling
        let epsilon_arc = Arc::epsilon(TropicalWeight::one(), 0);
        assert!(filter.filter_arc(0, 0, &epsilon_arc, &arc1)); // Epsilon matches any
        assert!(filter.filter_arc(0, 0, &arc1, &epsilon_arc)); // Any matches epsilon
    }

    #[test]
    fn test_compose_with_label_pair_lookahead() {
        // Create FSTs where lookahead can optimize composition
        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::new(1, 1, TropicalWeight::one(), s1));
        fst1.add_arc(s1, Arc::new(2, 2, TropicalWeight::one(), s2));

        let mut fst2 = VectorFst::<TropicalWeight>::new();
        let s0 = fst2.add_state();
        let s1 = fst2.add_state();
        let s2 = fst2.add_state();
        fst2.set_start(s0);
        fst2.set_final(s2, TropicalWeight::one());
        fst2.add_arc(s0, Arc::new(1, 3, TropicalWeight::one(), s1));
        fst2.add_arc(s1, Arc::new(2, 4, TropicalWeight::one(), s2));

        let filter = LabelPairLookaheadFilter::with_max_depth(1);
        let result: VectorFst<TropicalWeight> =
            compose_with_lookahead(&fst1, &fst2, filter).unwrap();
        assert!(result.num_states() > 0);
    }
}