arcweight 0.3.0

//! FST difference algorithm.
//!
//! Computes the language difference between two weighted finite-state acceptors,
//! creating an acceptor that recognizes strings accepted by the first but not the second.
//!
//! For acceptors $`A_1`$ and $`A_2`$:
//! $`L(A_1 - A_2) = L(A_1) \cap \overline{L(A_2)} = \{w : w \in L(A_1) \text{ and } w \notin L(A_2)\}`$
//!
//! # Implementation
//!
//! The algorithm constructs the complement of $`A_2`$ by:
//! 1. Determinizing $`A_2`$ to ensure one transition per symbol per state
//! 2. Making the FST complete by adding a sink state for missing transitions
//! 3. Flipping final states (final becomes non-final, non-final becomes final)
//!
//! The result is then computed as $`A_1 \cap \overline{A_2}`$ via intersection.
//!
//! # Complexity
//!
//! - **Time:** $`O(2^{|V_2|})`$ worst case for determinization, then $`O(|V_1| \times |V_2'|)`$ for intersection
//! - **Space:** $`O(|V_1| \times |V_2'|)`$ for state cross product after determinization
//!
//! where $`V_2'`$ is the determinized complement size.
//!
//! # References
//!
//! - John E. Hopcroft and Jeffrey D. Ullman. 1979. *Introduction to Automata Theory,
//!   Languages, and Computation*. Addison-Wesley, Reading, MA.
//! - Mehryar Mohri. 2009. Weighted automata algorithms. In *Handbook of Weighted
//!   Automata*, Manfred Droste, Werner Kuich, and Heiko Vogler (Eds.). Springer,
//!   Berlin, Heidelberg, 213-254.

use crate::algorithms::{determinize, intersect};
use crate::arc::Arc;
use crate::fst::{Fst, Label, MutableFst};
use crate::semiring::{DivisibleSemiring, Semiring};
use crate::{Error, Result};
use core::hash::Hash;
use std::collections::{HashMap, HashSet};

/// Computes the language difference between two finite-state acceptors.
///
/// Creates an acceptor recognizing strings accepted by the first but not the second:
/// $`L(A_1 - A_2) = \{w : w \in L(A_1) \text{ and } w \notin L(A_2)\}`$
///
/// # Arguments
///
/// * `fst1` - The first acceptor (minuend)
/// * `fst2` - The second acceptor (subtrahend)
///
/// # Type Parameters
///
/// * `W` - Weight type implementing [`DivisibleSemiring`] + [`Ord`] + [`Hash`]
/// * `F1` - First FST type implementing [`Fst<W>`]
/// * `F2` - Second FST type implementing [`Fst<W>`]
/// * `M` - Output FST type implementing [`MutableFst<W>`] and [`Default`]
///
/// # Returns
///
/// A new acceptor recognizing the language difference.
///
/// # Errors
///
/// Returns [`Error::Algorithm`] if:
/// - Either input FST is not an acceptor (input labels must equal output labels)
/// - The combined alphabet is empty (no non-epsilon labels)
/// - Determinization or intersection fails
///
/// # Examples
///
/// ```rust
/// use arcweight::prelude::*;
/// use arcweight::algorithms::difference;
///
/// // Acceptor 1: accepts "a" and "ab"
/// let mut acc1 = VectorFst::<TropicalWeight>::new();
/// let s0 = acc1.add_state();
/// let s1 = acc1.add_state();
/// let s2 = acc1.add_state();
/// acc1.set_start(s0);
/// acc1.set_final(s1, TropicalWeight::one()); // accepts "a"
/// acc1.set_final(s2, TropicalWeight::one()); // accepts "ab"
/// acc1.add_arc(s0, Arc::new(1, 1, TropicalWeight::one(), s1));
/// acc1.add_arc(s1, Arc::new(2, 2, TropicalWeight::one(), s2));
///
/// // Acceptor 2: accepts "ab"
/// let mut acc2 = VectorFst::<TropicalWeight>::new();
/// let s0 = acc2.add_state();
/// let s1 = acc2.add_state();
/// let s2 = acc2.add_state();
/// acc2.set_start(s0);
/// acc2.set_final(s2, TropicalWeight::one());
/// acc2.add_arc(s0, Arc::new(1, 1, TropicalWeight::one(), s1));
/// acc2.add_arc(s1, Arc::new(2, 2, TropicalWeight::one(), s2));
///
/// // Difference accepts "a" but not "ab"
/// let diff: VectorFst<TropicalWeight> = difference(&acc1, &acc2)?;
/// assert!(diff.start().is_some());
/// # Ok::<(), Box<dyn std::error::Error>>(())
/// ```
///
/// # References
///
/// - John E. Hopcroft and Jeffrey D. Ullman. 1979. *Introduction to Automata Theory,
///   Languages, and Computation*. Addison-Wesley, Reading, MA.
///
/// [`DivisibleSemiring`]: crate::semiring::DivisibleSemiring
pub fn difference<W, F1, F2, M>(fst1: &F1, fst2: &F2) -> Result<M>
where
    W: DivisibleSemiring + Hash + Clone + Ord + Eq,
    F1: Fst<W>,
    F2: Fst<W>,
    M: MutableFst<W> + Default,
{
    // Step 1: Validate inputs are acceptors
    validate_acceptor(fst1)?;
    validate_acceptor(fst2)?;

    // Step 2: Collect alphabet from both FSTs
    let alphabet = collect_alphabet(fst1, fst2)?;

    // Step 3: Build complement of fst2
    let complement_fst2: M = build_complement(fst2, &alphabet)?;

    // Step 4: Compute intersection of fst1 with complement(fst2)
    intersect(fst1, &complement_fst2)
}

/// Validate that an FST is an acceptor (input labels = output labels)
fn validate_acceptor<W: Semiring, F: Fst<W>>(fst: &F) -> Result<()> {
    for state in fst.states() {
        for arc in fst.arcs(state) {
            if arc.ilabel != arc.olabel {
                return Err(Error::Algorithm(
                    "FST must be an acceptor (input = output labels)".into(),
                ));
            }
        }
    }
    Ok(())
}

/// Collect the alphabet (all non-epsilon labels) used in both FSTs
fn collect_alphabet<W: Semiring, F1: Fst<W>, F2: Fst<W>>(
    fst1: &F1,
    fst2: &F2,
) -> Result<HashSet<Label>> {
    let mut alphabet = HashSet::new();

    // Collect labels from fst1
    for state in fst1.states() {
        for arc in fst1.arcs(state) {
            if arc.ilabel != 0 {
                // Skip epsilon
                alphabet.insert(arc.ilabel);
            }
        }
    }

    // Collect labels from fst2
    for state in fst2.states() {
        for arc in fst2.arcs(state) {
            if arc.ilabel != 0 {
                // Skip epsilon
                alphabet.insert(arc.ilabel);
            }
        }
    }

    if alphabet.is_empty() {
        return Err(Error::Algorithm(
            "Empty alphabet in difference operation".into(),
        ));
    }

    Ok(alphabet)
}

/// Build the complement acceptor for the given FST over the specified alphabet
fn build_complement<
    W: DivisibleSemiring + Clone + Ord + Hash + Eq,
    F: Fst<W>,
    M: MutableFst<W> + Default,
>(
    fst: &F,
    alphabet: &HashSet<Label>,
) -> Result<M> {
    // First, ensure the FST is deterministic and complete
    let det_fst: M = determinize(fst)?;
    let complete_fst: M = make_complete(&det_fst, alphabet)?;

    // Now build the complement by flipping final states
    let mut complement = M::default();

    // Copy all states
    let mut state_map = HashMap::new();
    for state in complete_fst.states() {
        let new_state = complement.add_state();
        state_map.insert(state, new_state);
    }

    // Set start state
    if let Some(start) = complete_fst.start() {
        if let Some(&new_start) = state_map.get(&start) {
            complement.set_start(new_start);
        }
    }

    // Copy all arcs
    for state in complete_fst.states() {
        if let Some(&new_state) = state_map.get(&state) {
            for arc in complete_fst.arcs(state) {
                if let Some(&new_nextstate) = state_map.get(&arc.nextstate) {
                    complement.add_arc(
                        new_state,
                        Arc::new(arc.ilabel, arc.olabel, arc.weight.clone(), new_nextstate),
                    );
                }
            }
        }
    }

    // Flip final states: non-final becomes final, final becomes non-final
    for state in complete_fst.states() {
        if let Some(&new_state) = state_map.get(&state) {
            if complete_fst.final_weight(state).is_none() {
                // Was non-final, make it final in complement
                complement.set_final(new_state, W::one());
            }
            // Was final, leave it non-final in complement (don't set final weight)
        }
    }

    Ok(complement)
}

/// Make an FST complete by adding transitions to a sink state for missing alphabet symbols
fn make_complete<
    W: DivisibleSemiring + Clone + Ord + Hash + Eq,
    F: Fst<W>,
    M: MutableFst<W> + Default,
>(
    fst: &F,
    alphabet: &HashSet<Label>,
) -> Result<M> {
    let mut complete = M::default();

    // Copy all original states
    let mut state_map = HashMap::new();
    for state in fst.states() {
        let new_state = complete.add_state();
        state_map.insert(state, new_state);
    }

    // Add sink state (non-final state that accepts everything)
    let sink_state = complete.add_state();

    // Set start state
    if let Some(start) = fst.start() {
        if let Some(&new_start) = state_map.get(&start) {
            complete.set_start(new_start);
        }
    }

    // Copy original arcs and final weights
    for state in fst.states() {
        if let Some(&new_state) = state_map.get(&state) {
            // Copy final weight
            if let Some(weight) = fst.final_weight(state) {
                complete.set_final(new_state, weight.clone());
            }

            // Copy arcs
            for arc in fst.arcs(state) {
                if let Some(&new_nextstate) = state_map.get(&arc.nextstate) {
                    complete.add_arc(
                        new_state,
                        Arc::new(arc.ilabel, arc.olabel, arc.weight.clone(), new_nextstate),
                    );
                }
            }
        }
    }

    // Add missing transitions to sink state
    for state in fst.states() {
        if let Some(&new_state) = state_map.get(&state) {
            // Collect existing outgoing labels for this state
            let mut existing_labels = HashSet::new();
            for arc in fst.arcs(state) {
                if arc.ilabel != 0 {
                    existing_labels.insert(arc.ilabel);
                }
            }

            // Add transitions to sink for missing labels
            for &label in alphabet {
                if !existing_labels.contains(&label) {
                    complete.add_arc(new_state, Arc::new(label, label, W::one(), sink_state));
                }
            }
        }
    }

    // Add self-loops on sink state for all alphabet symbols
    for &label in alphabet {
        complete.add_arc(sink_state, Arc::new(label, label, W::one(), sink_state));
    }

    Ok(complete)
}

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

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

        // Valid acceptor: input = output labels
        fst.add_arc(s0, Arc::new(1, 1, TropicalWeight::one(), s1));

        assert!(validate_acceptor(&fst).is_ok());
    }

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

        // Invalid: different input/output labels (transducer)
        fst.add_arc(s0, Arc::new(1, 2, TropicalWeight::one(), s1));

        assert!(validate_acceptor(&fst).is_err());
    }

    #[test]
    fn test_collect_alphabet_basic() {
        let mut fst1 = VectorFst::<TropicalWeight>::new();
        let mut fst2 = VectorFst::<TropicalWeight>::new();

        let s0 = fst1.add_state();
        let s1 = fst1.add_state();
        fst1.add_arc(s0, Arc::new(1, 1, TropicalWeight::one(), s1));
        fst1.add_arc(s1, Arc::new(2, 2, TropicalWeight::one(), s0));

        let s0 = fst2.add_state();
        let s1 = fst2.add_state();
        fst2.add_arc(s0, Arc::new(2, 2, TropicalWeight::one(), s1));
        fst2.add_arc(s1, Arc::new(3, 3, TropicalWeight::one(), s0));

        let alphabet = collect_alphabet(&fst1, &fst2).unwrap();
        assert_eq!(alphabet.len(), 3);
        assert!(alphabet.contains(&1));
        assert!(alphabet.contains(&2));
        assert!(alphabet.contains(&3));
    }

    #[test]
    fn test_collect_alphabet_with_epsilon() {
        let mut fst1 = VectorFst::<TropicalWeight>::new();
        let mut fst2 = VectorFst::<TropicalWeight>::new();

        let s0 = fst1.add_state();
        let s1 = fst1.add_state();
        fst1.add_arc(s0, Arc::new(0, 0, TropicalWeight::one(), s1)); // Epsilon
        fst1.add_arc(s1, Arc::new(1, 1, TropicalWeight::one(), s0));

        let s0 = fst2.add_state();
        let s1 = fst2.add_state();
        fst2.add_arc(s0, Arc::new(2, 2, TropicalWeight::one(), s1));

        let alphabet = collect_alphabet(&fst1, &fst2).unwrap();
        assert_eq!(alphabet.len(), 2);
        assert!(alphabet.contains(&1));
        assert!(alphabet.contains(&2));
        assert!(!alphabet.contains(&0)); // Epsilon should be excluded
    }

    #[test]
    fn test_collect_alphabet_empty() {
        let mut fst1 = VectorFst::<TropicalWeight>::new();
        let mut fst2 = VectorFst::<TropicalWeight>::new();

        // Only epsilon arcs
        let s0 = fst1.add_state();
        let s1 = fst1.add_state();
        fst1.add_arc(s0, Arc::new(0, 0, TropicalWeight::one(), s1));

        let _s0 = fst2.add_state();
        // No arcs in fst2

        let result = collect_alphabet(&fst1, &fst2);
        assert!(result.is_err()); // Should fail with empty alphabet
    }

    #[test]
    fn test_make_complete_basic() {
        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 mut alphabet = HashSet::new();
        alphabet.insert(1);
        alphabet.insert(2);

        let complete: VectorFst<TropicalWeight> = make_complete(&fst, &alphabet).unwrap();

        // Should have original states plus sink state
        assert_eq!(complete.num_states(), 3);
        assert!(complete.start().is_some());

        // Should have transitions for all alphabet symbols from all states
        assert!(complete.num_arcs_total() > fst.num_arcs_total());
    }

    #[test]
    fn test_build_complement_simple() {
        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 mut alphabet = HashSet::new();
        alphabet.insert(1);

        let complement: VectorFst<TropicalWeight> = build_complement(&fst, &alphabet).unwrap();

        // Complement should have more states (original + sink)
        assert!(complement.num_states() >= fst.num_states());
        assert!(complement.start().is_some());
    }

    #[test]
    fn test_difference_basic() {
        // FST1 accepts "a"
        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));

        // FST2 accepts "b"
        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, 2, TropicalWeight::one(), s1));

        // Difference should accept "a" (fst1 - fst2)
        let diff: VectorFst<TropicalWeight> = difference(&fst1, &fst2).unwrap();
        assert!(diff.start().is_some());
        assert!(diff.num_states() > 0);
    }

    #[test]
    fn test_difference_self() {
        // FST accepts "a"
        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));

        // Difference with itself should be empty
        let diff: VectorFst<TropicalWeight> = difference(&fst, &fst).unwrap();

        // Result should have no final states or should be empty
        let final_count = diff.states().filter(|&s| diff.is_final(s)).count();
        assert_eq!(final_count, 0);
    }

    #[test]
    fn test_difference_empty_fst() {
        let fst1 = VectorFst::<TropicalWeight>::new();
        let mut fst2 = VectorFst::<TropicalWeight>::new();
        let s0 = fst2.add_state();
        fst2.set_start(s0);
        fst2.set_final(s0, TropicalWeight::one());

        // Difference with empty FST should fail (empty alphabet)
        let result = difference::<TropicalWeight, _, _, VectorFst<TropicalWeight>>(&fst1, &fst2);
        assert!(result.is_err());
    }

    #[test]
    fn test_difference_non_acceptor() {
        // Create transducer (not acceptor)
        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)); // Different input/output

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

        // Should fail because fst1 is not an acceptor
        let result = difference::<TropicalWeight, _, _, VectorFst<TropicalWeight>>(&fst1, &fst2);
        assert!(result.is_err());
    }

    #[test]
    fn test_difference_overlapping_languages() {
        // FST1 accepts "a" and "ab"
        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(s1, TropicalWeight::one()); // accepts "a"
        fst1.set_final(s2, TropicalWeight::one()); // accepts "ab"
        fst1.add_arc(s0, Arc::new(1, 1, TropicalWeight::one(), s1)); // a
        fst1.add_arc(s1, Arc::new(2, 2, TropicalWeight::one(), s2)); // b

        // FST2 accepts "ab"
        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, 1, TropicalWeight::one(), s1)); // a
        fst2.add_arc(s1, Arc::new(2, 2, TropicalWeight::one(), s2)); // b

        // Difference should accept "a" but not "ab"
        let diff: VectorFst<TropicalWeight> = difference(&fst1, &fst2).unwrap();
        assert!(diff.start().is_some());
        assert!(diff.num_states() > 0);
    }

    #[test]
    fn test_difference_disjoint_languages() {
        // FST1 accepts "a"
        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));

        // FST2 accepts "b"
        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, 2, TropicalWeight::one(), s1));

        // Since languages are disjoint, difference should equal fst1
        let diff: VectorFst<TropicalWeight> = difference(&fst1, &fst2).unwrap();
        assert!(diff.start().is_some());
        assert!(diff.num_states() > 0);

        // Should have at least one final state
        let final_count = diff.states().filter(|&s| diff.is_final(s)).count();
        assert!(final_count > 0);
    }

    #[test]
    fn test_difference_complex_automata() {
        // FST1: accepts strings ending with "a"
        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)); // a (final)
        fst1.add_arc(s0, Arc::new(2, 2, TropicalWeight::one(), s0)); // b (stay)
        fst1.add_arc(s1, Arc::new(1, 1, TropicalWeight::one(), s1)); // a (final)
        fst1.add_arc(s1, Arc::new(2, 2, TropicalWeight::one(), s0)); // b (back to start)

        // FST2: accepts single "a"
        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, 1, TropicalWeight::one(), s1)); // a

        // Difference should accept strings ending with "a" except single "a"
        let diff: VectorFst<TropicalWeight> = difference(&fst1, &fst2).unwrap();
        assert!(diff.start().is_some());
        assert!(diff.num_states() > 0);
    }

    #[test]
    fn test_difference_preserves_weights() {
        // FST1 with weighted arcs
        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::new(2.0));
        fst1.add_arc(s0, Arc::new(1, 1, TropicalWeight::new(1.5), s1));

        // FST2: empty language (no final states)
        let mut fst2 = VectorFst::<TropicalWeight>::new();
        let s0 = fst2.add_state();
        fst2.set_start(s0);
        fst2.add_arc(s0, Arc::new(1, 1, TropicalWeight::one(), s0));
        // No final states - accepts nothing

        // Difference should preserve fst1 exactly
        let diff: VectorFst<TropicalWeight> = difference(&fst1, &fst2).unwrap();
        assert!(diff.start().is_some());

        // Check that some weights are preserved (structure may differ)
        let has_final = diff.states().any(|s| diff.is_final(s));
        assert!(has_final);
    }

    #[test]
    fn test_difference_single_state_acceptors() {
        // FST1: single state accepting empty string
        let mut fst1 = VectorFst::<TropicalWeight>::new();
        let s0 = fst1.add_state();
        fst1.set_start(s0);
        fst1.set_final(s0, TropicalWeight::one());
        fst1.add_arc(s0, Arc::new(1, 1, TropicalWeight::one(), s0)); // Self-loop on "a"

        // FST2: different single state
        let mut fst2 = VectorFst::<TropicalWeight>::new();
        let s0 = fst2.add_state();
        fst2.set_start(s0);
        fst2.set_final(s0, TropicalWeight::one());
        fst2.add_arc(s0, Arc::new(2, 2, TropicalWeight::one(), s0)); // Self-loop on "b"

        let diff: VectorFst<TropicalWeight> = difference(&fst1, &fst2).unwrap();
        assert!(diff.start().is_some());
    }

    #[test]
    fn test_difference_epsilon_handling() {
        // FST1 with epsilon transitions
        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(0, 0, TropicalWeight::one(), s1)); // epsilon
        fst1.add_arc(s1, Arc::new(1, 1, TropicalWeight::one(), s2)); // a

        // FST2: simple acceptor
        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, 2, TropicalWeight::one(), s1)); // b

        // Should handle epsilon transitions properly
        let diff: VectorFst<TropicalWeight> = difference(&fst1, &fst2).unwrap();
        assert!(diff.start().is_some());
    }
}