arcweight 0.3.0

//! FST pruning algorithm.
//!
//! Removes paths from weighted FSTs that exceed a weight threshold, producing a smaller
//! FST that preserves only the "best" paths according to the semiring ordering.
//!
//! Pruning is fundamental to beam search and other efficient search algorithms in speech
//! recognition, machine translation, and natural language processing. It enables trading
//! completeness for efficiency by discarding low-probability or high-cost paths early.
//!
//! # Semiring Requirements
//!
//! Pruning requires a **naturally ordered semiring** ([`NaturallyOrderedSemiring`]) for
//! meaningful weight comparison:
//! - Total ordering defines which paths are "worse" than the threshold
//! - Without natural ordering, the pruning concept is not well-defined
//!
//! **Supported semirings:**
//! - [`TropicalWeight`] - Natural ordering by cost (prune high-cost paths)
//! - [`LogWeight`] - Natural ordering by log probability
//!
//! **Unsupported semirings:**
//! - `ProbabilityWeight` - No natural ordering defined
//! - `BooleanWeight` - No meaningful weight comparison
//!
//! # Complexity
//!
//! - **Weight-based:** $`O(|V| \times |E|)`$ using Bellman-Ford shortest distances
//! - **Forward-backward:** $`O(|V| \times |E|)`$ with two Bellman-Ford passes
//! - **N-best paths:** $`O(|V| \times |E| \times \log N)`$ with priority queue
//!
//! # References
//!
//! - Mehryar Mohri, Fernando Pereira, and Michael Riley. 2002. Weighted finite-state
//!   transducers in speech recognition. *Computer Speech & Language* 16, 1 (2002),
//!   69-88. DOI: <https://doi.org/10.1006/csla.2001.0184>
//! - Mehryar Mohri. 2009. Weighted automata algorithms. In *Handbook of Weighted
//!   Automata*, Manfred Droste, Werner Kuich, and Heiko Vogler (Eds.). Springer,
//!   Berlin, Heidelberg, 213-254.
//!
//! [`NaturallyOrderedSemiring`]: crate::semiring::NaturallyOrderedSemiring
//! [`TropicalWeight`]: crate::semiring::TropicalWeight
//! [`LogWeight`]: crate::semiring::LogWeight

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

/// Configuration for the pruning algorithm.
///
/// Controls the pruning strategy and thresholds used to filter paths in the FST.
/// Multiple pruning criteria can be combined for fine-grained control.
///
/// # Examples
///
/// ```rust
/// use arcweight::algorithms::PruneConfig;
///
/// // Weight-based pruning with threshold of 5.0
/// let config = PruneConfig {
///     weight_threshold: 5.0,
///     ..Default::default()
/// };
///
/// // Forward-backward pruning for beam search
/// let beam_config = PruneConfig {
///     weight_threshold: 3.0,
///     use_forward_backward: true,
///     ..Default::default()
/// };
///
/// // N-best paths pruning
/// let nbest_config = PruneConfig {
///     npath: Some(10),
///     ..Default::default()
/// };
/// ```
#[derive(Debug, Clone)]
pub struct PruneConfig {
    /// Weight threshold for pruning
    pub weight_threshold: f64,
    /// Maximum number of states to keep
    pub state_threshold: Option<usize>,
    /// Number of shortest paths to keep
    pub npath: Option<usize>,
    /// Prune based on forward-backward weights
    pub use_forward_backward: bool,
    /// Delta for weight comparison
    pub delta: f64,
}

impl Default for PruneConfig {
    fn default() -> Self {
        Self {
            weight_threshold: f64::INFINITY,
            state_threshold: None,
            npath: None,
            use_forward_backward: false,
            delta: 1e-6,
        }
    }
}

/// State with priority for heap operations
#[derive(Debug, Clone)]
struct PriorityState<W> {
    state: StateId,
    weight: W,
}

impl<W: PartialOrd> PartialEq for PriorityState<W> {
    fn eq(&self, other: &Self) -> bool {
        self.weight == other.weight
    }
}

impl<W: PartialOrd> Eq for PriorityState<W> {}

impl<W: PartialOrd> Ord for PriorityState<W> {
    fn cmp(&self, other: &Self) -> Ordering {
        // Reverse order for min-heap behavior
        other
            .weight
            .partial_cmp(&self.weight)
            .unwrap_or(Ordering::Equal)
    }
}

impl<W: PartialOrd> PartialOrd for PriorityState<W> {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

/// Prunes an FST by removing paths and states that exceed weight thresholds.
///
/// Removes paths from the FST that have weights exceeding the specified threshold,
/// producing a smaller FST that accepts a subset of the original language.
///
/// **Correctness guarantee:** $`L(\text{prune}(T)) \subseteq L(T)`$ -- pruning only
/// removes paths, never adds them.
///
/// # Arguments
///
/// * `fst` - The input FST to prune
/// * `config` - Configuration specifying pruning thresholds and strategy
///
/// # Type Parameters
///
/// * `W` - Weight type implementing [`NaturallyOrderedSemiring`]
/// * `F` - Input FST type implementing [`Fst<W>`]
/// * `M` - Output FST type implementing [`MutableFst<W>`] and [`Default`]
///
/// # Returns
///
/// A new FST containing only paths within the specified weight thresholds.
///
/// # Errors
///
/// Returns [`Error::Algorithm`](crate::Error::Algorithm) if:
/// - The input FST has invalid structure
/// - Memory allocation fails during construction
///
/// # Examples
///
/// ```rust
/// use arcweight::prelude::*;
/// use arcweight::algorithms::{prune, PruneConfig};
///
/// let mut fst = VectorFst::<TropicalWeight>::new();
/// let s0 = fst.add_state();
/// let s1 = fst.add_state();
/// fst.set_start(s0);
/// fst.set_final(s1, TropicalWeight::one());
/// fst.add_arc(s0, Arc::new(1, 1, TropicalWeight::new(2.0), s1));
///
/// let config = PruneConfig {
///     weight_threshold: 5.0,
///     ..Default::default()
/// };
/// let pruned: VectorFst<TropicalWeight> = prune(&fst, config)?;
/// # Ok::<(), Box<dyn std::error::Error>>(())
/// ```
///
/// # References
///
/// - Mehryar Mohri, Fernando Pereira, and Michael Riley. 2002. Weighted finite-state
///   transducers in speech recognition. *Computer Speech & Language* 16, 1 (2002),
///   69-88.
///
/// [`NaturallyOrderedSemiring`]: crate::semiring::NaturallyOrderedSemiring
pub fn prune<W, F, M>(fst: &F, config: PruneConfig) -> Result<M>
where
    W: NaturallyOrderedSemiring,
    W::Value: Into<f64> + Copy,
    F: Fst<W>,
    M: MutableFst<W> + Default,
{
    // Handle empty FST case
    if fst.num_states() == 0 || fst.start().is_none() {
        return Ok(M::default());
    }

    // Choose pruning strategy based on configuration
    if config.use_forward_backward {
        prune_forward_backward(fst, &config)
    } else if let Some(npath) = config.npath {
        prune_nbest_paths(fst, npath, &config)
    } else {
        prune_by_weight(fst, &config)
    }
}

/// Prune using forward-backward algorithm
fn prune_forward_backward<W, F, M>(fst: &F, config: &PruneConfig) -> Result<M>
where
    W: NaturallyOrderedSemiring,
    W::Value: Into<f64> + Copy,
    F: Fst<W>,
    M: MutableFst<W> + Default,
{
    let forward_weights = compute_forward_weights(fst)?;
    let backward_weights = compute_backward_weights(fst)?;

    // Find best total weight
    let mut best_weight = None;
    if let Some(start) = fst.start() {
        if let Some(backward) = backward_weights.get(&start) {
            best_weight = Some(backward.clone());
        }
    }

    let mut result = M::default();
    let mut state_map = FxHashMap::default();

    // First pass: select states based on forward-backward weights
    let zero_weight = W::zero();
    for state in fst.states() {
        let forward = forward_weights.get(&state).unwrap_or(&zero_weight);
        let backward = backward_weights.get(&state).unwrap_or(&zero_weight);

        if *forward == W::zero() || *backward == W::zero() {
            continue; // Skip unreachable states
        }

        // Compute total weight through this state
        let total = forward.times(backward);

        // Check if state should be kept
        if should_keep_weight(&total, &best_weight, config) {
            let new_state = result.add_state();
            state_map.insert(state, new_state);

            // Copy final weight if applicable
            if let Some(weight) = fst.final_weight(state) {
                result.set_final(new_state, weight.clone());
            }
        }
    }

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

    // Second pass: copy arcs between selected states
    for (&old_state, &new_state) in &state_map {
        for arc in fst.arcs(old_state) {
            if let Some(&new_nextstate) = state_map.get(&arc.nextstate) {
                // Check arc weight against threshold
                let arc_total = if let (Some(forward), Some(backward)) = (
                    forward_weights.get(&old_state),
                    backward_weights.get(&arc.nextstate),
                ) {
                    forward.times(&arc.weight).times(backward)
                } else {
                    continue;
                };

                if should_keep_weight(&arc_total, &best_weight, config) {
                    result.add_arc(
                        new_state,
                        Arc::new(arc.ilabel, arc.olabel, arc.weight.clone(), new_nextstate),
                    );
                }
            }
        }
    }

    apply_state_threshold(result, config)
}

/// Prune keeping only n-best paths
fn prune_nbest_paths<W, F, M>(fst: &F, npath: usize, config: &PruneConfig) -> Result<M>
where
    W: NaturallyOrderedSemiring,
    W::Value: Into<f64> + Copy,
    F: Fst<W>,
    M: MutableFst<W> + Default,
{
    if npath == 0 {
        return Ok(M::default());
    }

    let start = fst
        .start()
        .ok_or_else(|| crate::Error::Algorithm("FST has no start state".into()))?;

    // Find n-best paths using priority queue
    let mut heap = BinaryHeap::new();
    let mut best_weights: FxHashMap<StateId, Vec<W>> = FxHashMap::default();

    // Initialize with start state
    heap.push(PriorityState {
        state: start,
        weight: W::one(),
    });
    best_weights.insert(start, vec![W::one()]);

    // Process states in best-first order
    while let Some(PriorityState { state, weight }) = heap.pop() {
        // Check if this is a final state
        if let Some(final_weight) = fst.final_weight(state) {
            let _total = weight.times(final_weight);
            // Track as a complete path
        }

        // Explore outgoing arcs
        for arc in fst.arcs(state) {
            let new_weight = weight.times(&arc.weight);

            // Update best weights for next state
            let weights = best_weights.entry(arc.nextstate).or_default();

            // Keep only n-best weights
            if weights.len() < npath {
                weights.push(new_weight.clone());
                weights.sort();
                if weights.len() > npath {
                    weights.truncate(npath);
                }

                heap.push(PriorityState {
                    state: arc.nextstate,
                    weight: new_weight,
                });
            } else if weights.last().is_some_and(|w| new_weight < *w) {
                // Replace worst weight if this is better
                weights[npath - 1] = new_weight.clone();
                weights.sort();

                heap.push(PriorityState {
                    state: arc.nextstate,
                    weight: new_weight,
                });
            }
        }
    }

    // Build result FST with selected paths
    build_nbest_fst(fst, &best_weights, npath, config)
}

/// Simple weight-based pruning
fn prune_by_weight<W, F, M>(fst: &F, config: &PruneConfig) -> Result<M>
where
    W: NaturallyOrderedSemiring,
    W::Value: Into<f64> + Copy,
    F: Fst<W>,
    M: MutableFst<W> + Default,
{
    let mut result = M::default();
    let mut state_map = FxHashMap::default();

    // Compute shortest distances from start
    let distances = compute_shortest_distances(fst)?;

    // First pass: select states within threshold
    for state in fst.states() {
        if let Some(distance) = distances.get(&state) {
            if convert_weight_to_f64(distance) <= config.weight_threshold {
                let new_state = result.add_state();
                state_map.insert(state, new_state);

                if let Some(weight) = fst.final_weight(state) {
                    result.set_final(new_state, weight.clone());
                }
            }
        }
    }

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

    // Second pass: copy arcs with weight filtering
    for (&old_state, &new_state) in &state_map {
        if let Some(state_distance) = distances.get(&old_state) {
            for arc in fst.arcs(old_state) {
                // Check if arc leads to selected state
                if let Some(&new_nextstate) = state_map.get(&arc.nextstate) {
                    // Compute path weight through this arc
                    let arc_distance = state_distance.times(&arc.weight);

                    if convert_weight_to_f64(&arc_distance) <= config.weight_threshold {
                        result.add_arc(
                            new_state,
                            Arc::new(arc.ilabel, arc.olabel, arc.weight.clone(), new_nextstate),
                        );
                    }
                }
            }
        }
    }

    apply_state_threshold(result, config)
}

/// Compute forward weights (shortest distance from start)
///
/// Uses Bellman-Ford relaxation to compute shortest distances from the start state.
///
/// # Complexity
///
/// - **Acyclic graphs:** O(|V| + |E|)
/// - **Cyclic graphs:** O(|V| × |E|) worst case, with iteration limit
///
/// # Algorithm
///
/// 1. Initialize start state with weight one (identity)
/// 2. Relaxation: for each arc (u, v) with weight w:
///    - If dist\[u\] ⊗ w < dist\[v\], update dist\[v\]
/// 3. Continue until no updates (or iteration limit reached)
///
/// # Termination
///
/// - For acyclic graphs: guaranteed to terminate in |V| iterations
/// - For cyclic graphs: may have improving cycles, so we limit iterations to |V|
///
/// # Notes
///
/// For graphs with negative (improving) cycles in the weight space, this may not
/// find true shortest distances but will terminate and return approximate values.
fn compute_forward_weights<W, F>(fst: &F) -> Result<FxHashMap<StateId, W>>
where
    W: NaturallyOrderedSemiring,
    F: Fst<W>,
{
    let mut weights = FxHashMap::default();
    let mut queue = VecDeque::new();
    let mut in_queue = FxHashSet::default();

    if let Some(start) = fst.start() {
        weights.insert(start, W::one());
        queue.push_back(start);
        in_queue.insert(start);
    }

    // Bellman-Ford style relaxation with iteration limit
    // Bellman-Ford requires at most n-1 iterations for graphs without negative cycles.
    // Using n iterations allows detecting if we're still making progress (indicating cycles).
    let num_states = fst.num_states();
    let max_iterations = num_states;
    let mut iterations = 0;

    while let Some(state) = queue.pop_front() {
        in_queue.remove(&state);
        iterations += 1;

        if iterations > max_iterations {
            // Prevent infinite loops on cyclic graphs
            break;
        }

        let state_weight = weights[&state].clone();

        for arc in fst.arcs(state) {
            let new_weight = state_weight.times(&arc.weight);

            let updated = match weights.get(&arc.nextstate) {
                None => {
                    weights.insert(arc.nextstate, new_weight);
                    true
                }
                Some(old_weight) => {
                    if new_weight < *old_weight {
                        weights.insert(arc.nextstate, new_weight);
                        true
                    } else {
                        false
                    }
                }
            };

            if updated && !in_queue.contains(&arc.nextstate) {
                queue.push_back(arc.nextstate);
                in_queue.insert(arc.nextstate);
            }
        }
    }

    Ok(weights)
}

/// Compute backward weights (shortest distance to final states)
///
/// Computes shortest distances from each state to any final state using
/// backward relaxation through reversed arcs.
///
/// # Complexity
///
/// - **Build reverse index:** O(|E|)
/// - **Relaxation:** O(|V| × |E|) worst case with iteration limit
/// - **Total:** O(|V| × |E|)
///
/// # Algorithm
///
/// 1. Build reverse arc index: for each arc (u, v), store (u, w) at v
/// 2. Initialize final states with their final weights
/// 3. Backward relaxation: for each state v, update predecessors u:
///    - dist\[u\] = dist\[u\] ⊕ (w ⊗ dist\[v\])
///
/// # Termination
///
/// Similar to forward weights, uses iteration limit for cyclic graphs.
fn compute_backward_weights<W, F>(fst: &F) -> Result<FxHashMap<StateId, W>>
where
    W: NaturallyOrderedSemiring,
    F: Fst<W>,
{
    let mut weights = FxHashMap::default();
    let mut reverse_arcs: FxHashMap<StateId, Vec<(StateId, W)>> = FxHashMap::default();

    // Build reverse arc index
    for state in fst.states() {
        for arc in fst.arcs(state) {
            reverse_arcs
                .entry(arc.nextstate)
                .or_default()
                .push((state, arc.weight.clone()));
        }
    }

    // Initialize with final states
    let mut queue = VecDeque::new();
    let mut in_queue = FxHashSet::default();
    for state in fst.states() {
        if let Some(final_weight) = fst.final_weight(state) {
            weights.insert(state, final_weight.clone());
            queue.push_back(state);
            in_queue.insert(state);
        }
    }

    // Backward relaxation with iteration limit
    // Bellman-Ford requires at most n-1 iterations; use n to detect cycles.
    let num_states = fst.num_states();
    let max_iterations = num_states;
    let mut iterations = 0;

    while let Some(state) = queue.pop_front() {
        in_queue.remove(&state);
        iterations += 1;

        if iterations > max_iterations {
            // Prevent infinite loops on cyclic graphs
            break;
        }

        let state_weight = weights[&state].clone();

        if let Some(predecessors) = reverse_arcs.get(&state) {
            for (prev_state, arc_weight) in predecessors {
                let new_weight = arc_weight.times(&state_weight);

                let updated = match weights.get(prev_state) {
                    None => {
                        weights.insert(*prev_state, new_weight);
                        true
                    }
                    Some(old_weight) => {
                        let combined = old_weight.plus(&new_weight);
                        if combined != *old_weight {
                            weights.insert(*prev_state, combined);
                            true
                        } else {
                            false
                        }
                    }
                };

                if updated && !in_queue.contains(prev_state) {
                    queue.push_back(*prev_state);
                    in_queue.insert(*prev_state);
                }
            }
        }
    }

    Ok(weights)
}

/// Compute shortest distances from start state
fn compute_shortest_distances<W, F>(fst: &F) -> Result<FxHashMap<StateId, W>>
where
    W: NaturallyOrderedSemiring,
    F: Fst<W>,
{
    compute_forward_weights(fst)
}

/// Build FST from n-best paths
fn build_nbest_fst<W, F, M>(
    fst: &F,
    best_weights: &FxHashMap<StateId, Vec<W>>,
    _npath: usize,
    _config: &PruneConfig,
) -> Result<M>
where
    W: NaturallyOrderedSemiring,
    F: Fst<W>,
    M: MutableFst<W> + Default,
{
    let mut result = M::default();
    let mut state_map = FxHashMap::default();

    // Create states that appear in best paths
    for &state in best_weights.keys() {
        let new_state = result.add_state();
        state_map.insert(state, new_state);

        if let Some(weight) = fst.final_weight(state) {
            result.set_final(new_state, weight.clone());
        }
    }

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

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

    Ok(result)
}

/// Apply state threshold by selecting best states
///
/// This function implements state limiting by ranking states based on their
/// contribution to the best paths and keeping only the top `threshold` states.
///
/// # Algorithm
///
/// 1. Compute forward weights (shortest distance from start to each state)
/// 2. Compute backward weights (shortest distance from each state to final states)
/// 3. Compute importance for each state: forward_weight × backward_weight
///    (represents the weight of the best path going through this state)
/// 4. Rank states by importance (lower is better for naturally ordered semirings)
/// 5. Keep the `threshold` states with lowest importance values
/// 6. Build a new FST containing only the selected states and their connecting arcs
///
/// # Complexity
///
/// - **Time:** O(V × E) for forward/backward weight computation + O(V log V) for sorting
/// - **Space:** O(V + E) for weight maps and result FST
///
/// # Notes
///
/// - The start state is always included in the selection (if it has paths to final states)
///   even if it's not among the most important states, to ensure the FST remains functional
/// - States unreachable from start or unable to reach final states are excluded
/// - The resulting FST may have fewer than `threshold` states if the original
///   FST has fewer reachable states
/// - After the start state is included, remaining states are selected by importance
/// - Ties in importance are broken by state ID (lower IDs preferred for stability)
fn apply_state_threshold<W, M>(fst: M, config: &PruneConfig) -> Result<M>
where
    W: NaturallyOrderedSemiring,
    W::Value: Into<f64> + Copy,
    M: MutableFst<W> + Default,
{
    let threshold = match config.state_threshold {
        Some(t) => t,
        None => return Ok(fst),
    };

    if fst.num_states() <= threshold {
        return Ok(fst);
    }

    // Need start state for forward weights
    let start = match fst.start() {
        Some(s) => s,
        None => return Ok(fst), // No start state, return as-is
    };

    // Compute forward weights (shortest distance from start to each state)
    let forward_weights = compute_forward_weights(&fst)?;

    // Compute backward weights (shortest distance from each state to final states)
    let backward_weights = compute_backward_weights(&fst)?;

    // Compute importance for each state: forward × backward
    // This represents the weight of the best path going through this state
    // Lower importance = better (for tropical/log semirings)
    let mut state_importance: Vec<(StateId, f64)> = Vec::new();

    for state in fst.states() {
        let forward = forward_weights.get(&state);
        let backward = backward_weights.get(&state);

        match (forward, backward) {
            (Some(fw), Some(bw)) => {
                // State is on at least one complete path (start -> state -> final)
                // Skip if either weight is semiring zero (infinity in tropical/log = unreachable)
                // Note: We use Semiring::is_zero, not num_traits::Zero::is_zero, because
                // TropicalWeight's semiring zero is infinity, not 0.0
                if Semiring::is_zero(fw) || Semiring::is_zero(bw) {
                    continue;
                }

                let importance = fw.times(bw);
                let importance_val = convert_weight_to_f64(&importance);

                // Only include states with finite importance
                if importance_val.is_finite() {
                    state_importance.push((state, importance_val));
                }
            }
            _ => {
                // State is not on any complete path (unreachable or dead-end)
                // Skip it - it contributes nothing to the language
            }
        }
    }

    // Handle edge case: no states on complete paths
    if state_importance.is_empty() {
        return Ok(M::default());
    }

    // Sort by importance (lower is better), with state ID as tiebreaker for stability
    state_importance.sort_by(|a, b| match a.1.partial_cmp(&b.1) {
        Some(std::cmp::Ordering::Equal) | None => a.0.cmp(&b.0),
        Some(ord) => ord,
    });

    // Select top `threshold` states (those with lowest importance values)
    // Always include the start state to ensure the FST remains functional
    let mut selected_states: FxHashSet<StateId> = FxHashSet::default();

    // First, add the start state if it has finite importance
    let start_in_importance = state_importance.iter().any(|(s, _)| *s == start);
    if !start_in_importance {
        // Start state has infinite importance (no path to final), return empty FST
        return Ok(M::default());
    }
    selected_states.insert(start);

    // Then add remaining states by importance until we reach threshold
    for (state, _) in &state_importance {
        if selected_states.len() >= threshold {
            break;
        }
        selected_states.insert(*state);
    }

    // Build new FST with selected states
    let mut result = M::default();
    let mut state_map: FxHashMap<StateId, StateId> = FxHashMap::default();

    // Create new states for all selected states
    for &old_state in &selected_states {
        let new_state = result.add_state();
        state_map.insert(old_state, new_state);
    }

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

    // Copy final weights and arcs for selected states
    for &old_state in &selected_states {
        let new_state = match state_map.get(&old_state) {
            Some(&s) => s,
            None => continue,
        };

        // Copy final weight if present
        if let Some(weight) = fst.final_weight(old_state) {
            result.set_final(new_state, weight.clone());
        }

        // Copy arcs that connect to other selected states
        for arc in fst.arcs(old_state) {
            if let Some(&new_nextstate) = state_map.get(&arc.nextstate) {
                result.add_arc(
                    new_state,
                    Arc::new(arc.ilabel, arc.olabel, arc.weight.clone(), new_nextstate),
                );
            }
        }
    }

    Ok(result)
}

/// Check if weight should be kept based on pruning criteria
fn should_keep_weight<W>(weight: &W, best: &Option<W>, config: &PruneConfig) -> bool
where
    W: NaturallyOrderedSemiring,
    W::Value: Into<f64> + Copy,
{
    // Check against absolute threshold
    if convert_weight_to_f64(weight) > config.weight_threshold {
        return false;
    }

    // Check against best weight with beam
    if let Some(best_weight) = best {
        let weight_val = convert_weight_to_f64(weight);
        let best_val = convert_weight_to_f64(best_weight);

        if weight_val > best_val + config.weight_threshold {
            return false;
        }
    }

    true
}

/// Convert weight to f64 for threshold comparison
///
/// Uses the `Semiring::value()` method to extract the underlying numeric value
/// and converts it to f64 for comparison with threshold values.
fn convert_weight_to_f64<W>(weight: &W) -> f64
where
    W: Semiring,
    W::Value: Into<f64> + Copy,
{
    (*weight.value()).into()
}

/// Compute reachable states from a given start state
#[allow(dead_code)]
fn compute_reachable_states<F: Fst<W>, W: Semiring>(fst: &F, start: StateId) -> FxHashSet<StateId> {
    let mut reachable = FxHashSet::default();
    let mut stack = vec![start];

    while let Some(state) = stack.pop() {
        if reachable.insert(state) {
            for arc in fst.arcs(state) {
                stack.push(arc.nextstate);
            }
        }
    }

    reachable
}

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

    #[test]
    fn test_prune_config_default() {
        let config = PruneConfig::default();
        assert_eq!(config.weight_threshold, f64::INFINITY);
        assert_eq!(config.state_threshold, None);
        assert_eq!(config.npath, None);
        assert!(!config.use_forward_backward);
    }

    #[test]
    fn test_prune_config_custom() {
        let config = PruneConfig {
            weight_threshold: 5.0,
            state_threshold: Some(100),
            npath: Some(10),
            use_forward_backward: true,
            delta: 1e-8,
        };
        assert_eq!(config.weight_threshold, 5.0);
        assert_eq!(config.state_threshold, Some(100));
        assert_eq!(config.npath, Some(10));
        assert!(config.use_forward_backward);
    }

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

        fst.set_start(s0);
        fst.set_final(s2, TropicalWeight::one());
        fst.add_arc(s0, Arc::new(1, 1, TropicalWeight::new(1.0), s1));
        fst.add_arc(s1, Arc::new(2, 2, TropicalWeight::new(2.0), s2));

        let config = PruneConfig {
            weight_threshold: 10.0,
            ..Default::default()
        };
        let pruned: VectorFst<TropicalWeight> = prune(&fst, config).unwrap();

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

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

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

        // Low-cost and high-cost paths
        fst.add_arc(s0, Arc::new(1, 1, TropicalWeight::new(0.5), s1));
        fst.add_arc(s0, Arc::new(2, 2, TropicalWeight::new(5.0), s2));

        let config = PruneConfig {
            weight_threshold: 3.0,
            ..Default::default()
        };
        let pruned: VectorFst<TropicalWeight> = prune(&fst, config).unwrap();

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

        // Should keep low-cost path
        if let Some(start) = pruned.start() {
            let reachable = compute_reachable_states(&pruned, start);
            assert!(reachable.iter().any(|&s| pruned.is_final(s)));
        }
    }

    #[test]
    fn test_prune_empty_fst() {
        let fst = VectorFst::<TropicalWeight>::new();
        let config = PruneConfig::default();
        let pruned: VectorFst<TropicalWeight> = prune(&fst, config).unwrap();

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

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

        let config = PruneConfig {
            weight_threshold: 5.0,
            ..Default::default()
        };
        let pruned: VectorFst<TropicalWeight> = prune(&fst, config).unwrap();

        assert_eq!(pruned.num_states(), 1);
        assert!(pruned.start().is_some());
    }

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

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

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

        let config = PruneConfig {
            weight_threshold: 5.0,
            use_forward_backward: true,
            ..Default::default()
        };
        let pruned: VectorFst<TropicalWeight> = prune(&fst, config).unwrap();

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

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

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

        // Multiple paths with different costs
        fst.add_arc(s0, Arc::new(1, 1, TropicalWeight::new(1.0), s1));
        fst.add_arc(s1, Arc::new(2, 2, TropicalWeight::new(0.5), s3));

        fst.add_arc(s0, Arc::new(3, 3, TropicalWeight::new(2.0), s2));
        fst.add_arc(s2, Arc::new(4, 4, TropicalWeight::new(0.5), s3));

        let config = PruneConfig {
            npath: Some(1),
            ..Default::default()
        };
        let pruned: VectorFst<TropicalWeight> = prune(&fst, config).unwrap();

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

    #[test]
    fn test_convert_weight_to_f64() {
        let w1 = TropicalWeight::new(3.5);
        let val1 = convert_weight_to_f64(&w1);
        assert!((val1 - 3.5).abs() < 1e-6);

        let w2 = TropicalWeight::zero();
        let val2 = convert_weight_to_f64(&w2);
        assert_eq!(val2, f64::INFINITY);
    }

    #[test]
    fn test_priority_state() {
        let ps1 = PriorityState {
            state: 0,
            weight: TropicalWeight::new(1.0),
        };
        let ps2 = PriorityState {
            state: 1,
            weight: TropicalWeight::new(2.0),
        };

        // ps1 should have higher priority (lower weight)
        assert!(ps1 > ps2);
    }

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

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

        for i in 0..9 {
            fst.add_arc(
                states[i],
                Arc::new(
                    (i + 1) as u32,
                    (i + 1) as u32,
                    TropicalWeight::new(0.1),
                    states[i + 1],
                ),
            );
        }

        let config = PruneConfig {
            state_threshold: Some(5),
            ..Default::default()
        };
        let pruned: VectorFst<TropicalWeight> = prune(&fst, config).unwrap();

        assert!(pruned.num_states() > 0);
    }

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

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

        // Diamond-shaped graph
        fst.add_arc(s0, Arc::new(1, 1, TropicalWeight::new(1.0), s1));
        fst.add_arc(s0, Arc::new(2, 2, TropicalWeight::new(2.0), s2));
        fst.add_arc(s1, Arc::new(3, 3, TropicalWeight::new(1.0), s3));
        fst.add_arc(s2, Arc::new(4, 4, TropicalWeight::new(1.0), s3));

        let config = PruneConfig {
            weight_threshold: 2.5,
            ..Default::default()
        };
        let pruned: VectorFst<TropicalWeight> = prune(&fst, config).unwrap();

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

        // Should keep the better path (through s1)
        if let Some(start) = pruned.start() {
            let reachable = compute_reachable_states(&pruned, start);
            assert!(reachable.iter().any(|&s| pruned.is_final(s)));
        }
    }

    #[test]
    fn test_state_threshold_limits_states() {
        // Create FST with 10 states in a chain
        let mut fst = VectorFst::<TropicalWeight>::new();
        let states: Vec<_> = (0..10).map(|_| fst.add_state()).collect();

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

        // Chain: s0 -> s1 -> s2 -> ... -> s9
        for i in 0..9 {
            fst.add_arc(
                states[i],
                Arc::new(
                    (i + 1) as u32,
                    (i + 1) as u32,
                    TropicalWeight::new(1.0),
                    states[i + 1],
                ),
            );
        }

        // Limit to 5 states
        let config = PruneConfig {
            state_threshold: Some(5),
            ..Default::default()
        };
        let pruned: VectorFst<TropicalWeight> = prune(&fst, config).unwrap();

        // Should have at most 5 states
        assert!(
            pruned.num_states() <= 5,
            "Expected at most 5 states, got {}",
            pruned.num_states()
        );
        assert!(pruned.start().is_some());
    }

    #[test]
    fn test_state_threshold_keeps_best_states() {
        // Create FST with diamond pattern - two paths with different costs
        let mut fst = VectorFst::<TropicalWeight>::new();
        let s0 = fst.add_state(); // start
        let s1 = fst.add_state(); // good path intermediate
        let s2 = fst.add_state(); // bad path intermediate
        let s3 = fst.add_state(); // final

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

        // Good path: s0 -> s1 -> s3 (weight 1.0 + 1.0 = 2.0)
        fst.add_arc(s0, Arc::new(1, 1, TropicalWeight::new(1.0), s1));
        fst.add_arc(s1, Arc::new(2, 2, TropicalWeight::new(1.0), s3));

        // Bad path: s0 -> s2 -> s3 (weight 5.0 + 5.0 = 10.0)
        fst.add_arc(s0, Arc::new(3, 3, TropicalWeight::new(5.0), s2));
        fst.add_arc(s2, Arc::new(4, 4, TropicalWeight::new(5.0), s3));

        // Limit to 3 states (should exclude s2 as it's on the worse path)
        let config = PruneConfig {
            state_threshold: Some(3),
            ..Default::default()
        };
        let pruned: VectorFst<TropicalWeight> = prune(&fst, config).unwrap();

        // Should have exactly 3 states
        assert_eq!(pruned.num_states(), 3);

        // Should have path to final
        assert!(pruned.start().is_some());
        let start = pruned.start().unwrap();
        let reachable = compute_reachable_states(&pruned, start);
        assert!(
            reachable.iter().any(|&s| pruned.is_final(s)),
            "Pruned FST should still have a path to a final state"
        );
    }

    #[test]
    fn test_state_threshold_preserves_final_weights() {
        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::new(0.5));

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

        let config = PruneConfig {
            state_threshold: Some(5), // More than we have
            ..Default::default()
        };
        let pruned: VectorFst<TropicalWeight> = prune(&fst, config).unwrap();

        // Should preserve the structure
        assert_eq!(pruned.num_states(), 2);

        // Should preserve final weight
        let mut has_final = false;
        for state in pruned.states() {
            if let Some(w) = pruned.final_weight(state) {
                has_final = true;
                assert_eq!(*w.value(), 0.5);
            }
        }
        assert!(has_final, "Should have preserved final state with weight");
    }

    #[test]
    fn test_state_threshold_excludes_dead_ends() {
        // Create FST with a dead-end state (no path to final)
        let mut fst = VectorFst::<TropicalWeight>::new();
        let s0 = fst.add_state(); // start
        let s1 = fst.add_state(); // leads to final
        let s2 = fst.add_state(); // dead end
        let s3 = fst.add_state(); // final

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

        // Good path: s0 -> s1 -> s3
        fst.add_arc(s0, Arc::new(1, 1, TropicalWeight::new(1.0), s1));
        fst.add_arc(s1, Arc::new(2, 2, TropicalWeight::new(1.0), s3));

        // Dead end: s0 -> s2 (s2 has no path to any final state)
        fst.add_arc(s0, Arc::new(3, 3, TropicalWeight::new(0.1), s2));

        // Limit to 3 states
        let config = PruneConfig {
            state_threshold: Some(3),
            ..Default::default()
        };
        let pruned: VectorFst<TropicalWeight> = prune(&fst, config).unwrap();

        // Dead-end state s2 should be excluded even though its forward weight is low
        // because it has infinite backward weight (no path to final)
        assert_eq!(pruned.num_states(), 3);

        // All states in result should be on complete paths
        let start = pruned.start().unwrap();
        let reachable = compute_reachable_states(&pruned, start);
        assert!(reachable.iter().any(|&s| pruned.is_final(s)));
    }

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

        fst.set_start(s0);
        fst.set_final(s2, TropicalWeight::new(0.5)); // Better final
        fst.set_final(s3, TropicalWeight::new(2.0)); // Worse final

        // Path to s2: weight 1.0
        fst.add_arc(s0, Arc::new(1, 1, TropicalWeight::new(1.0), s1));
        fst.add_arc(s1, Arc::new(2, 2, TropicalWeight::new(0.5), s2));

        // Path to s3: weight 3.0
        fst.add_arc(s1, Arc::new(3, 3, TropicalWeight::new(2.5), s3));

        // Limit to 3 states - should prefer path to better final
        let config = PruneConfig {
            state_threshold: Some(3),
            ..Default::default()
        };
        let pruned: VectorFst<TropicalWeight> = prune(&fst, config).unwrap();

        assert!(pruned.num_states() <= 3);
        assert!(pruned.start().is_some());
    }

    #[test]
    fn test_state_threshold_no_start() {
        let mut fst = VectorFst::<TropicalWeight>::new();
        let s0 = fst.add_state();
        // No start state set
        fst.set_final(s0, TropicalWeight::one());

        let config = PruneConfig {
            state_threshold: Some(1),
            ..Default::default()
        };
        let pruned: VectorFst<TropicalWeight> = prune(&fst, config).unwrap();

        // Should return empty FST since there's no start state for weight-based pruning
        // (the default prune_by_weight handles no-start case)
        assert!(pruned.start().is_none() || pruned.num_states() == 0);
    }

    #[test]
    fn test_state_threshold_all_states_on_best_path() {
        // When all states contribute equally to the best path
        let mut fst = VectorFst::<TropicalWeight>::new();
        let s0 = fst.add_state();
        let s1 = fst.add_state();
        let s2 = fst.add_state();

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

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

        // Request 2 states from a 3-state chain
        let config = PruneConfig {
            state_threshold: Some(2),
            ..Default::default()
        };
        let pruned: VectorFst<TropicalWeight> = prune(&fst, config).unwrap();

        // Should have at most 2 states
        assert!(pruned.num_states() <= 2);
    }

    #[test]
    fn test_state_threshold_below_current_count() {
        // Test when threshold is below current state count
        let mut fst = VectorFst::<TropicalWeight>::new();
        let states: Vec<_> = (0..20).map(|_| fst.add_state()).collect();

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

        // Create a chain
        for i in 0..19 {
            fst.add_arc(
                states[i],
                Arc::new(1, 1, TropicalWeight::new(0.5), states[i + 1]),
            );
        }

        // Request only 5 states from a 20-state chain
        let config = PruneConfig {
            state_threshold: Some(5),
            ..Default::default()
        };
        let pruned: VectorFst<TropicalWeight> = prune(&fst, config).unwrap();

        // Should have at most 5 states
        assert!(
            pruned.num_states() <= 5,
            "Expected at most 5 states, got {}",
            pruned.num_states()
        );
    }

    #[test]
    fn test_state_threshold_equal_to_current() {
        // Test when threshold equals current state count
        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 config = PruneConfig {
            state_threshold: Some(2), // Exactly the number of states
            ..Default::default()
        };
        let pruned: VectorFst<TropicalWeight> = prune(&fst, config).unwrap();

        // Should preserve all states
        assert_eq!(pruned.num_states(), 2);
    }

    #[test]
    fn test_state_threshold_larger_than_current() {
        // Test when threshold is larger than current state count
        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 config = PruneConfig {
            state_threshold: Some(100), // Much larger than state count
            ..Default::default()
        };
        let pruned: VectorFst<TropicalWeight> = prune(&fst, config).unwrap();

        // Should preserve all states
        assert_eq!(pruned.num_states(), 2);
    }
}