arcweight 0.3.0

//! Parallel algorithm implementations for multi-core FST processing.
//!
//! This module provides parallelized versions of core FST algorithms
//! for improved performance on multi-core systems using work-stealing
//! and frontier-based parallelism strategies.
//!
//! # Overview
//!
//! Parallel FST algorithms exploit the inherent parallelism in graph-based
//! computations by processing independent states or state pairs concurrently.
//! This module implements several parallelization strategies:
//!
//! - **Frontier-based parallelism:** Process all states at the same BFS depth simultaneously
//! - **Work-stealing:** Dynamic load balancing across worker threads
//! - **Parallel relaxation:** Concurrent edge relaxation for shortest path algorithms
//! - **Delta-stepping:** Bucket-based parallelization for shortest distances
//!
//! # Feature Gate
//!
//! These algorithms require the `parallel` feature (enabled by default).
//! When disabled, all functions fall back to sequential implementations.
//!
//! # Algorithms
//!
//! - [`compose_parallel`] - Parallel FST composition using frontier-based parallelism
//! - [`map_weights_parallel`] - Parallel weight transformation across all arcs
//! - [`parallel_state_map`] - Generic parallel state processing
//! - [`collect_arcs_parallel`] - Parallel arc collection for batch operations
//! - [`shortest_distance_parallel`] - Parallel shortest distance with Bellman-Ford
//! - [`delta_stepping_shortest_distance`] - Delta-stepping parallel shortest paths
//! - [`determinize_parallel`] - Parallel subset construction with work-stealing
//! - [`minimize_parallel`] - Parallel partition refinement for minimization
//!
//! # Complexity
//!
//! For P processors and an FST with |V| states and |E| arcs:
//!
//! | Algorithm | Time | Space |
//! |-----------|------|-------|
//! | `compose_parallel` | $`O((V_1 V_2 + E_1 E_2) / P)`$ | $`O(V_1 V_2)`$ |
//! | `map_weights_parallel` | $`O(E / P)`$ | $`O(V)`$ |
//! | `shortest_distance_parallel` | $`O(V \cdot E / P)`$ | $`O(V)`$ |
//! | `delta_stepping_shortest_distance` | $`O((V + E/\delta + L \cdot D) / P)`$ | $`O(V + E/\delta)`$ |
//! | `determinize_parallel` | $`O(2^V / P)`$ typical | $`O(2^V)`$ |
//! | `minimize_parallel` | $`O(n \log n / P)`$ | $`O(n)`$ |
//!
//! # Performance Guidelines
//!
//! - **Small FSTs (<1000 states):** Sequential algorithms may be faster due to
//!   parallelization overhead. Use the non-parallel variants.
//! - **Medium FSTs (1000-100000 states):** Parallel algorithms provide good speedups.
//! - **Large FSTs (>100000 states):** Parallel algorithms essential for practical runtimes.
//! - **Core count:** Best performance with 4+ physical cores.
//!
//! # Thread Safety
//!
//! All parallel algorithms use thread-safe synchronization:
//! - `RwLock` for read-heavy shared state
//! - `Mutex` for write-heavy shared state
//! - Atomic operations for counters and flags
//! - Lock-free algorithms where possible
//!
//! # References
//!
//! - Meyer, U., and Sanders, P. (2003). Delta-stepping: A parallelizable shortest path
//!   algorithm. *Journal of Algorithms*, 49(1), 114-152.
//!   <https://doi.org/10.1016/S0196-6774(03)00076-2>
//! - Allauzen, C., Riley, M., Schalkwyk, J., Skut, W., and Mohri, M. (2007).
//!   OpenFst: A general and efficient weighted finite-state transducer library.
//!   In *International Conference on Implementation and Application of Automata*
//!   (pp. 11-23). Springer. <https://doi.org/10.1007/978-3-540-76336-9_3>
//! - Leiserson, C. E., and Schardl, T. B. (2010). A work-efficient parallel
//!   breadth-first search algorithm. In *Proceedings of the Twenty-Second Annual
//!   ACM Symposium on Parallelism in Algorithms and Architectures* (pp. 303-314).
//!   <https://doi.org/10.1145/1810479.1810534>

#[cfg(feature = "parallel")]
use rayon::prelude::*;

use crate::arc::Arc;
use crate::fst::{Fst, MutableFst, StateId};
use crate::semiring::Semiring;
#[cfg(feature = "parallel")]
use crate::Error;
use crate::Result;

#[cfg(feature = "parallel")]
use rustc_hash::{FxHashMap, FxHashSet};
#[cfg(feature = "parallel")]
use std::sync::{Arc as StdArc, Mutex, RwLock};

/// Parallel weight mapping across all arcs
///
/// Applies a transformation function to all arc weights in parallel.
/// This is particularly useful for bulk weight updates.
///
/// # Complexity
///
/// - **Time:** O(E / P) where E = total arcs, P = number of processors
/// - **Space:** O(V) for state-level synchronization
///
/// # Examples
///
/// ```rust
/// use arcweight::prelude::*;
/// use arcweight::algorithms::parallel::map_weights_parallel;
///
/// 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(1.0), s1));
///
/// // Scale all weights
/// let scaled: VectorFst<TropicalWeight> = map_weights_parallel(&fst, |w| w.times(&TropicalWeight::new(2.0))).unwrap();
/// ```
#[cfg(feature = "parallel")]
pub fn map_weights_parallel<W, F, M, MapFn>(fst: &F, map_fn: MapFn) -> Result<M>
where
    W: Semiring + Send + Sync,
    F: Fst<W> + Sync,
    M: MutableFst<W> + Default + Send,
    MapFn: Fn(&W) -> W + Send + Sync,
{
    let num_states = fst.num_states();

    // Collect all state data in parallel
    let state_data: Vec<_> = (0..num_states as StateId)
        .into_par_iter()
        .map(|state| {
            let final_weight = fst.final_weight(state).map(&map_fn);
            let arcs: Vec<_> = fst
                .arcs(state)
                .map(|arc| Arc::new(arc.ilabel, arc.olabel, map_fn(&arc.weight), arc.nextstate))
                .collect();
            (final_weight, arcs)
        })
        .collect();

    // Build result FST
    let mut result = M::default();

    for _ in 0..num_states {
        result.add_state();
    }

    if let Some(start) = fst.start() {
        result.set_start(start);
    }

    for (state, (final_weight, arcs)) in state_data.into_iter().enumerate() {
        let state_id = state as StateId;

        if let Some(w) = final_weight {
            result.set_final(state_id, w);
        }

        for arc in arcs {
            result.add_arc(state_id, arc);
        }
    }

    Ok(result)
}

/// Non-parallel fallback for map_weights_parallel
#[cfg(not(feature = "parallel"))]
pub fn map_weights_parallel<W, F, M, MapFn>(fst: &F, map_fn: MapFn) -> Result<M>
where
    W: Semiring,
    F: Fst<W>,
    M: MutableFst<W> + Default,
    MapFn: Fn(&W) -> W,
{
    let num_states = fst.num_states();
    let mut result = M::default();

    for _ in 0..num_states {
        result.add_state();
    }

    if let Some(start) = fst.start() {
        result.set_start(start);
    }

    for state in 0..num_states as StateId {
        if let Some(w) = fst.final_weight(state) {
            result.set_final(state, map_fn(w));
        }

        for arc in fst.arcs(state) {
            result.add_arc(
                state,
                Arc::new(arc.ilabel, arc.olabel, map_fn(&arc.weight), arc.nextstate),
            );
        }
    }

    Ok(result)
}

/// Parallel state processing for FST algorithms
///
/// Processes each state independently in parallel, collecting results
/// that can be merged into a final FST.
///
/// # Type Parameters
///
/// - `W`: Semiring weight type
/// - `F`: Input FST type
/// - `R`: Result type from processing each state
/// - `ProcessFn`: Function to process each state
#[cfg(feature = "parallel")]
pub fn parallel_state_map<W, F, R, ProcessFn>(fst: &F, process: ProcessFn) -> Vec<R>
where
    W: Semiring + Send + Sync,
    F: Fst<W> + Sync,
    R: Send,
    ProcessFn: Fn(StateId, &F) -> R + Send + Sync,
{
    let num_states = fst.num_states();

    (0..num_states as StateId)
        .into_par_iter()
        .map(|state| process(state, fst))
        .collect()
}

/// Non-parallel fallback
#[cfg(not(feature = "parallel"))]
pub fn parallel_state_map<W, F, R, ProcessFn>(fst: &F, process: ProcessFn) -> Vec<R>
where
    W: Semiring,
    F: Fst<W>,
    ProcessFn: Fn(StateId, &F) -> R,
{
    let num_states = fst.num_states();

    (0..num_states as StateId)
        .map(|state| process(state, fst))
        .collect()
}

/// Parallel arc collection for batch processing
///
/// Collects all arcs from an FST in parallel, suitable for bulk operations.
#[cfg(feature = "parallel")]
pub fn collect_arcs_parallel<W, F>(fst: &F) -> Vec<(StateId, Vec<Arc<W>>)>
where
    W: Semiring + Send + Sync,
    F: Fst<W> + Sync,
{
    let num_states = fst.num_states();

    (0..num_states as StateId)
        .into_par_iter()
        .map(|state| {
            let arcs: Vec<_> = fst.arcs(state).collect();
            (state, arcs)
        })
        .collect()
}

/// Non-parallel fallback
#[cfg(not(feature = "parallel"))]
pub fn collect_arcs_parallel<W, F>(fst: &F) -> Vec<(StateId, Vec<Arc<W>>)>
where
    W: Semiring,
    F: Fst<W>,
{
    let num_states = fst.num_states();

    (0..num_states as StateId)
        .map(|state| {
            let arcs: Vec<_> = fst.arcs(state).collect();
            (state, arcs)
        })
        .collect()
}

/// Compute shortest distances in parallel for independent states
///
/// For FSTs where states can be processed independently (e.g., after
/// topological sorting), this computes distances in parallel.
#[cfg(feature = "parallel")]
pub fn parallel_relax_arcs<W>(distances: &[W], arcs_by_state: &[(StateId, Vec<Arc<W>>)]) -> Vec<W>
where
    W: Semiring + Send + Sync,
{
    let result: Vec<Mutex<W>> = distances.iter().map(|w| Mutex::new(w.clone())).collect();

    arcs_by_state.par_iter().for_each(|(state, arcs)| {
        let state_dist = distances[*state as usize].clone();

        if !<W as num_traits::Zero>::is_zero(&state_dist) {
            for arc in arcs {
                let next_weight = state_dist.times(&arc.weight);
                let mut target = result[arc.nextstate as usize].lock().unwrap();
                *target = target.plus(&next_weight);
            }
        }
    });

    result
        .into_iter()
        .map(|m| m.into_inner().unwrap())
        .collect()
}

/// Parallel FST composition using frontier-based parallelism
///
/// This algorithm processes state pairs at the same BFS depth in parallel,
/// providing significant speedups for large FSTs on multi-core systems.
///
/// # Complexity
///
/// - **Time:** O((V₁V₂ + E₁E₂) / P) where P = number of processors
/// - **Space:** O(V₁V₂) for state pair storage with lock-free access
///
/// # Algorithm
///
/// Uses frontier-based parallelism:
/// 1. Initialize frontier with start state pair
/// 2. While frontier is not empty:
///    - Process all state pairs in frontier in parallel
///    - Collect newly discovered state pairs
///    - Add new pairs to next frontier
/// 3. Build result FST from discovered structure
///
/// # Examples
///
/// ```rust
/// use arcweight::prelude::*;
/// use arcweight::algorithms::parallel::compose_parallel;
///
/// let mut fst1 = VectorFst::<TropicalWeight>::new();
/// let s0 = fst1.add_state();
/// let s1 = fst1.add_state();
/// fst1.set_start(s0);
/// fst1.set_final(s1, TropicalWeight::one());
/// fst1.add_arc(s0, Arc::new(1, 2, TropicalWeight::new(1.0), s1));
///
/// let mut fst2 = VectorFst::<TropicalWeight>::new();
/// let t0 = fst2.add_state();
/// let t1 = fst2.add_state();
/// fst2.set_start(t0);
/// fst2.set_final(t1, TropicalWeight::one());
/// fst2.add_arc(t0, Arc::new(2, 3, TropicalWeight::new(2.0), t1));
///
/// let composed: VectorFst<TropicalWeight> = compose_parallel(&fst1, &fst2).unwrap();
/// assert!(composed.num_states() > 0);
/// ```
///
/// # Performance Notes
///
/// - For small FSTs (<1000 states), use sequential `compose_default` instead
/// - Best performance with 4+ cores and large FSTs
/// - Memory overhead is similar to sequential composition
///
/// See the main documentation above for details.
#[cfg(feature = "parallel")]
pub fn compose_parallel<W, F1, F2, M>(fst1: &F1, fst2: &F2) -> Result<M>
where
    W: Semiring + Send + Sync,
    F1: Fst<W> + Sync,
    F2: Fst<W> + Sync,
    M: MutableFst<W> + Default + Send,
{
    let start1 = fst1
        .start()
        .ok_or_else(|| Error::Algorithm("First FST has no start state".into()))?;
    let start2 = fst2
        .start()
        .ok_or_else(|| Error::Algorithm("Second FST has no start state".into()))?;

    // Thread-safe state map: (s1, s2) -> composed state id
    let state_map: StdArc<RwLock<FxHashMap<(StateId, StateId), StateId>>> =
        StdArc::new(RwLock::new(FxHashMap::default()));

    // Counter for assigning new state IDs
    let next_state_id: StdArc<Mutex<StateId>> = StdArc::new(Mutex::new(0));

    // Collected arc data: (source_state, arc)
    let arc_data: StdArc<Mutex<Vec<(StateId, Arc<W>)>>> = StdArc::new(Mutex::new(Vec::new()));

    // Final state data: (state, weight)
    let final_data: StdArc<Mutex<Vec<(StateId, W)>>> = StdArc::new(Mutex::new(Vec::new()));

    // Register start state
    {
        let mut map = state_map.write().unwrap();
        let mut counter = next_state_id.lock().unwrap();
        map.insert((start1, start2), *counter);
        *counter += 1;
    }

    // Initialize frontier with start state pair
    let mut frontier: Vec<(StateId, StateId, StateId)> = vec![(start1, start2, 0)];
    let mut visited: FxHashSet<(StateId, StateId)> = FxHashSet::default();
    visited.insert((start1, start2));

    // Result type for parallel processing: (new_states, arcs, final_weight, current_state_id)
    type ProcessResult<W> = (
        Vec<(StateId, StateId, StateId)>, // new state pairs with their IDs
        Vec<(StateId, Arc<W>)>,           // arcs to add
        Option<W>,                        // final weight if applicable
        StateId,                          // current composed state ID
    );

    // BFS with parallel frontier processing
    while !frontier.is_empty() {
        // Process all state pairs in current frontier in parallel
        let frontier_results: Vec<ProcessResult<W>> = frontier
            .par_iter()
            .map(|(s1, s2, current_state)| {
                let mut local_new_states: Vec<(StateId, StateId, StateId)> = Vec::new();
                let mut local_arcs: Vec<(StateId, Arc<W>)> = Vec::new();
                let mut local_final: Option<W> = None;

                // Handle final states
                if let (Some(w1), Some(w2)) = (fst1.final_weight(*s1), fst2.final_weight(*s2)) {
                    local_final = Some(w1.times(w2));
                }

                // Process arc pairs
                for arc1 in fst1.arcs(*s1) {
                    for arc2 in fst2.arcs(*s2) {
                        // Match when output of FST1 equals input of FST2
                        if arc1.olabel == arc2.ilabel {
                            let next_pair = (arc1.nextstate, arc2.nextstate);

                            // Check if state pair exists (read lock)
                            let existing = {
                                let map = state_map.read().unwrap();
                                map.get(&next_pair).copied()
                            };

                            let next_state = match existing {
                                Some(state) => state,
                                None => {
                                    // Need to create new state (write lock)
                                    let mut map = state_map.write().unwrap();
                                    // Double-check after acquiring write lock
                                    if let Some(&state) = map.get(&next_pair) {
                                        state
                                    } else {
                                        let mut counter = next_state_id.lock().unwrap();
                                        let new_state = *counter;
                                        *counter += 1;
                                        map.insert(next_pair, new_state);
                                        local_new_states.push((
                                            next_pair.0,
                                            next_pair.1,
                                            new_state,
                                        ));
                                        new_state
                                    }
                                }
                            };

                            // Create composed arc
                            let composed_weight = arc1.weight.times(&arc2.weight);
                            local_arcs.push((
                                *current_state,
                                Arc::new(arc1.ilabel, arc2.olabel, composed_weight, next_state),
                            ));
                        }
                    }

                    // Handle epsilon output from FST1 (olabel = 0)
                    if arc1.olabel == 0 {
                        let next_pair = (arc1.nextstate, *s2);

                        let existing = {
                            let map = state_map.read().unwrap();
                            map.get(&next_pair).copied()
                        };

                        let next_state = match existing {
                            Some(state) => state,
                            None => {
                                let mut map = state_map.write().unwrap();
                                if let Some(&state) = map.get(&next_pair) {
                                    state
                                } else {
                                    let mut counter = next_state_id.lock().unwrap();
                                    let new_state = *counter;
                                    *counter += 1;
                                    map.insert(next_pair, new_state);
                                    local_new_states.push((next_pair.0, next_pair.1, new_state));
                                    new_state
                                }
                            }
                        };

                        local_arcs.push((
                            *current_state,
                            Arc::new(arc1.ilabel, 0, arc1.weight.clone(), next_state),
                        ));
                    }
                }

                // Handle epsilon input in FST2 (ilabel = 0)
                for arc2 in fst2.arcs(*s2) {
                    if arc2.ilabel == 0 {
                        let next_pair = (*s1, arc2.nextstate);

                        let existing = {
                            let map = state_map.read().unwrap();
                            map.get(&next_pair).copied()
                        };

                        let next_state = match existing {
                            Some(state) => state,
                            None => {
                                let mut map = state_map.write().unwrap();
                                if let Some(&state) = map.get(&next_pair) {
                                    state
                                } else {
                                    let mut counter = next_state_id.lock().unwrap();
                                    let new_state = *counter;
                                    *counter += 1;
                                    map.insert(next_pair, new_state);
                                    local_new_states.push((next_pair.0, next_pair.1, new_state));
                                    new_state
                                }
                            }
                        };

                        local_arcs.push((
                            *current_state,
                            Arc::new(0, arc2.olabel, arc2.weight.clone(), next_state),
                        ));
                    }
                }

                (local_new_states, local_arcs, local_final, *current_state)
            })
            .collect();

        // Collect results and prepare next frontier
        let mut next_frontier = Vec::new();

        for (new_states, arcs, final_w, current_state_id) in frontier_results {
            // Add arcs
            if !arcs.is_empty() {
                arc_data.lock().unwrap().extend(arcs);
            }

            // Add final weight
            if let Some(w) = final_w {
                final_data.lock().unwrap().push((current_state_id, w));
            }

            // Add newly discovered states to next frontier
            for (s1, s2, state_id) in new_states {
                if !visited.contains(&(s1, s2)) {
                    visited.insert((s1, s2));
                    next_frontier.push((s1, s2, state_id));
                }
            }
        }

        // Also check state_map for any states created by other threads that we haven't visited
        let map = state_map.read().unwrap();
        for (pair, &state_id) in map.iter() {
            if !visited.contains(pair) {
                visited.insert(*pair);
                next_frontier.push((pair.0, pair.1, state_id));
            }
        }
        drop(map);

        frontier = next_frontier;
    }

    // Build the result FST
    let num_states = *next_state_id.lock().unwrap() as usize;
    let mut result = M::default();

    // Add all states
    for _ in 0..num_states {
        result.add_state();
    }

    // Set start state
    result.set_start(0);

    // Add final weights
    for (state, weight) in final_data.lock().unwrap().drain(..) {
        if let Some(existing) = result.final_weight(state) {
            result.set_final(state, existing.plus(&weight));
        } else {
            result.set_final(state, weight);
        }
    }

    // Add all arcs
    for (source, arc) in arc_data.lock().unwrap().drain(..) {
        result.add_arc(source, arc);
    }

    Ok(result)
}

/// Non-parallel fallback for compose_parallel
///
/// Uses sequential composition when the parallel feature is disabled.
#[cfg(not(feature = "parallel"))]
pub fn compose_parallel<W, F1, F2, M>(fst1: &F1, fst2: &F2) -> Result<M>
where
    W: Semiring,
    F1: Fst<W>,
    F2: Fst<W>,
    M: MutableFst<W> + Default,
{
    // Fall back to sequential compose
    crate::algorithms::compose_default(fst1, fst2)
}

/// Parallel shortest distance computation
///
/// Computes shortest distances from the start state to all reachable states
/// using parallel relaxation with Bellman-Ford style convergence checking.
///
/// # Complexity
///
/// - **Time:** O(V * E / P) worst case where P = number of processors
/// - **Space:** O(V) for distance storage
///
/// # Examples
///
/// ```rust
/// use arcweight::prelude::*;
/// use arcweight::algorithms::parallel::shortest_distance_parallel;
///
/// 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 distances = shortest_distance_parallel(&fst).unwrap();
/// ```
#[cfg(feature = "parallel")]
pub fn shortest_distance_parallel<W, F>(fst: &F) -> Result<Vec<W>>
where
    W: Semiring + Send + Sync,
    F: Fst<W> + Sync,
{
    let num_states = fst.num_states();
    if num_states == 0 {
        return Ok(Vec::new());
    }

    let start = match fst.start() {
        Some(s) => s,
        None => return Ok(vec![W::zero(); num_states]),
    };

    // Initialize distances
    let distances: Vec<Mutex<W>> = (0..num_states)
        .map(|i| {
            if i as StateId == start {
                Mutex::new(W::one())
            } else {
                Mutex::new(W::zero())
            }
        })
        .collect();

    // Use Bellman-Ford style relaxation with convergence check
    // For idempotent semirings (like tropical), this converges quickly
    let max_iterations = num_states;
    let changed = StdArc::new(std::sync::atomic::AtomicBool::new(true));

    for _ in 0..max_iterations {
        // Use Acquire to synchronize with Release stores from previous iteration
        if !changed.load(std::sync::atomic::Ordering::Acquire) {
            break;
        }
        // Use Release to ensure all distance updates are visible before next iteration checks
        changed.store(false, std::sync::atomic::Ordering::Release);

        // Parallel relaxation of all edges
        (0..num_states as StateId)
            .into_par_iter()
            .for_each(|state| {
                let state_dist = distances[state as usize].lock().unwrap().clone();

                if !<W as num_traits::Zero>::is_zero(&state_dist) {
                    for arc in fst.arcs(state) {
                        let new_dist = state_dist.times(&arc.weight);
                        let mut target = distances[arc.nextstate as usize].lock().unwrap();
                        let combined = target.plus(&new_dist);
                        if combined != *target {
                            *target = combined;
                            // Use Release to ensure distance update is visible before setting flag
                            changed.store(true, std::sync::atomic::Ordering::Release);
                        }
                    }
                }
            });
    }

    // Extract final distances
    Ok(distances
        .into_iter()
        .map(|m| m.into_inner().unwrap())
        .collect())
}

/// Non-parallel fallback for shortest_distance_parallel
#[cfg(not(feature = "parallel"))]
pub fn shortest_distance_parallel<W, F>(fst: &F) -> Result<Vec<W>>
where
    W: Semiring,
    F: Fst<W>,
{
    crate::algorithms::shortest_distance(fst)
}

/// Delta-stepping parallel shortest path algorithm
///
/// This algorithm processes nodes in "buckets" by distance, enabling parallel
/// processing within each bucket. It's particularly effective for large graphs
/// with positive edge weights.
///
/// # Algorithm
///
/// Delta-stepping partitions nodes by their tentative distance into buckets:
/// - Bucket i contains nodes with distance in [i*δ, (i+1)*δ)
/// - "Light" edges (weight < δ) are processed within buckets (may cause relaxation within same bucket)
/// - "Heavy" edges (weight ≥ δ) are processed between buckets
///
/// # Complexity
///
/// - **Time:** O((V + E/δ + L·D) / P) where:
///   - δ = bucket width (delta parameter)
///   - L = maximum number of light edges relaxed per node
///   - D = diameter of the graph
///   - P = number of processors
/// - **Space:** O(V + E/δ) for bucket storage
///
/// # Parameters
///
/// - `fst`: Input FST
/// - `delta`: Bucket width (smaller = more parallelism, larger = fewer iterations)
///
/// # Examples
///
/// ```rust
/// use arcweight::prelude::*;
/// use arcweight::algorithms::parallel::delta_stepping_shortest_distance;
///
/// 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));
///
/// // Use delta of 1.0 for good parallelism
/// let distances = delta_stepping_shortest_distance(&fst, 1.0).unwrap();
/// assert_eq!(*distances[0].value(), 0.0); // Start
/// assert_eq!(*distances[1].value(), 1.0); // Via first arc
/// assert_eq!(*distances[2].value(), 3.0); // Via both arcs
/// ```
///
/// # References
///
/// Meyer, U., & Sanders, P. (2003). "Δ-stepping: A parallelizable shortest path algorithm."
/// Journal of Algorithms, 49(1), 114-152.
#[cfg(feature = "parallel")]
pub fn delta_stepping_shortest_distance<W, F>(fst: &F, delta: f64) -> Result<Vec<W>>
where
    W: Semiring + Send + Sync + Clone + PartialOrd,
    W::Value: Into<f64> + Copy,
    F: Fst<W> + Sync,
{
    use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};

    let num_states = fst.num_states();
    if num_states == 0 {
        return Ok(Vec::new());
    }

    let start = match fst.start() {
        Some(s) => s,
        None => return Ok(vec![W::zero(); num_states]),
    };

    // Initialize distances (using Mutex for thread-safe updates)
    let distances: Vec<Mutex<W>> = (0..num_states)
        .map(|i| {
            if i as StateId == start {
                Mutex::new(W::one())
            } else {
                Mutex::new(W::zero())
            }
        })
        .collect();

    // Pre-compute light and heavy edges for each state
    // Light edges: weight < delta (in appropriate semiring metric)
    // Heavy edges: weight >= delta
    let light_edges: Vec<Vec<(StateId, W)>> = (0..num_states)
        .map(|state| {
            fst.arcs(state as StateId)
                .filter(|arc| is_light_edge(&arc.weight, delta))
                .map(|arc| (arc.nextstate, arc.weight.clone()))
                .collect()
        })
        .collect();

    let heavy_edges: Vec<Vec<(StateId, W)>> = (0..num_states)
        .map(|state| {
            fst.arcs(state as StateId)
                .filter(|arc| !is_light_edge(&arc.weight, delta))
                .map(|arc| (arc.nextstate, arc.weight.clone()))
                .collect()
        })
        .collect();

    // Buckets: array of sets of nodes
    // Using a simpler approach with atomic flags for bucket membership
    let max_buckets = ((num_states as f64) * 2.0 / delta).ceil() as usize + 1;
    let buckets: Vec<Mutex<FxHashSet<StateId>>> = (0..max_buckets)
        .map(|_| Mutex::new(FxHashSet::default()))
        .collect();

    // Add start node to bucket 0
    buckets[0].lock().unwrap().insert(start);

    // Current bucket index
    let current_bucket = AtomicUsize::new(0);
    let any_updates = AtomicBool::new(true);

    // Process buckets
    while any_updates.load(Ordering::Relaxed) {
        any_updates.store(false, Ordering::Relaxed);

        let bucket_idx = current_bucket.load(Ordering::Relaxed);

        // Get nodes in current bucket
        let mut current_nodes: Vec<StateId> = Vec::new();
        {
            let mut bucket = buckets[bucket_idx].lock().unwrap();
            current_nodes.extend(bucket.drain());
        }

        if current_nodes.is_empty() {
            // Move to next non-empty bucket
            let next_bucket = buckets
                .iter()
                .enumerate()
                .skip(bucket_idx + 1)
                .take(max_buckets - bucket_idx - 1)
                .find(|(_, b)| !b.lock().unwrap().is_empty())
                .map(|(i, _)| i);

            if let Some(next_idx) = next_bucket {
                current_bucket.store(next_idx, Ordering::Relaxed);
                any_updates.store(true, Ordering::Relaxed);
                continue;
            } else {
                break;
            }
        }

        any_updates.store(true, Ordering::Relaxed);

        // Phase 1: Relax light edges (may add nodes back to current bucket)
        let mut light_iterations = 0;
        let max_light_iterations = num_states; // Prevent infinite loops

        loop {
            if current_nodes.is_empty() || light_iterations >= max_light_iterations {
                break;
            }
            light_iterations += 1;

            // Parallel relaxation of light edges
            let updates: Vec<Vec<(StateId, W)>> = current_nodes
                .par_iter()
                .map(|&node| {
                    let node_dist = distances[node as usize].lock().unwrap().clone();
                    let mut local_updates = Vec::new();

                    if !<W as num_traits::Zero>::is_zero(&node_dist) {
                        for (next_state, weight) in &light_edges[node as usize] {
                            let new_dist = node_dist.times(weight);
                            local_updates.push((*next_state, new_dist));
                        }
                    }

                    local_updates
                })
                .collect();

            // Apply updates and collect nodes that improved
            current_nodes.clear();
            for node_updates in updates {
                for (next_state, new_dist) in node_updates {
                    let mut target = distances[next_state as usize].lock().unwrap();
                    let combined = target.plus(&new_dist);
                    if combined != *target {
                        *target = combined;
                        // Re-add to current bucket for further processing
                        current_nodes.push(next_state);
                    }
                }
            }

            // Deduplicate
            current_nodes.sort();
            current_nodes.dedup();
        }

        // Phase 2: Relax heavy edges (add nodes to future buckets)
        // Get all nodes that were processed (including those added during light phase)
        let processed_nodes: Vec<StateId> = {
            let mut nodes: Vec<StateId> = (0..num_states as StateId)
                .filter(|&s| {
                    let dist = distances[s as usize].lock().unwrap();
                    !<W as num_traits::Zero>::is_zero(&dist)
                })
                .collect();
            nodes.sort();
            nodes.dedup();
            nodes
        };

        // Parallel relaxation of heavy edges
        let heavy_updates: Vec<Vec<(StateId, W)>> = processed_nodes
            .par_iter()
            .map(|&node| {
                let node_dist = distances[node as usize].lock().unwrap().clone();
                let mut local_updates = Vec::new();

                if !<W as num_traits::Zero>::is_zero(&node_dist) {
                    for (next_state, weight) in &heavy_edges[node as usize] {
                        let new_dist = node_dist.times(weight);
                        local_updates.push((*next_state, new_dist));
                    }
                }

                local_updates
            })
            .collect();

        // Apply heavy edge updates
        for node_updates in heavy_updates {
            for (next_state, new_dist) in node_updates {
                let mut target = distances[next_state as usize].lock().unwrap();
                let combined = target.plus(&new_dist);
                if combined != *target {
                    // Compute bucket for the combined distance
                    let actual_bucket = compute_bucket_index(&combined, delta, max_buckets);
                    *target = combined;
                    // Add to appropriate future bucket if it's ahead of current
                    if actual_bucket > bucket_idx && actual_bucket < max_buckets {
                        buckets[actual_bucket].lock().unwrap().insert(next_state);
                    }
                }
            }
        }

        // Move to next bucket
        current_bucket.fetch_add(1, Ordering::Relaxed);
    }

    // Extract final distances
    Ok(distances
        .into_iter()
        .map(|m| m.into_inner().unwrap())
        .collect())
}

/// Helper: Check if an edge is "light" (weight < delta in appropriate metric)
///
/// For numeric semirings like TropicalWeight, extracts the value and compares with delta.
/// Light edges (weight < delta) are processed more frequently in the delta-stepping algorithm.
#[cfg(feature = "parallel")]
fn is_light_edge<W: Semiring>(weight: &W, delta: f64) -> bool
where
    W::Value: Into<f64> + Copy,
{
    let weight_val: f64 = (*weight.value()).into();
    weight_val < delta
}

/// Helper: Compute bucket index for a given distance
///
/// For numeric semirings, bucket = floor(dist / delta), which distributes
/// nodes across buckets based on their tentative distance.
#[cfg(feature = "parallel")]
fn compute_bucket_index<W: Semiring>(dist: &W, delta: f64, max_buckets: usize) -> usize
where
    W::Value: Into<f64> + Copy,
{
    if <W as num_traits::Zero>::is_zero(dist) {
        max_buckets // Unreachable nodes go to last bucket
    } else {
        let dist_val: f64 = (*dist.value()).into();
        let bucket = (dist_val / delta).floor() as usize;
        bucket.min(max_buckets - 1)
    }
}

/// Non-parallel fallback for delta_stepping_shortest_distance
#[cfg(not(feature = "parallel"))]
pub fn delta_stepping_shortest_distance<W, F>(fst: &F, _delta: f64) -> Result<Vec<W>>
where
    W: Semiring + Clone + PartialOrd,
    W::Value: Into<f64> + Copy,
    F: Fst<W>,
{
    crate::algorithms::shortest_distance(fst)
}

/// Parallel determinization with work-stealing
///
/// This algorithm uses a work-stealing approach for parallelizing the subset
/// construction algorithm. Each worker thread processes subsets from a shared
/// work queue, enabling dynamic load balancing.
///
/// # Algorithm
///
/// 1. Initialize work queue with start subset
/// 2. Workers steal subsets from the queue
/// 3. Each worker:
///    - Processes a subset to find outgoing transitions
///    - Creates new subsets for unvisited destinations
///    - Adds new subsets to the work queue
/// 4. Workers use atomic operations to avoid conflicts
/// 5. Result FST is assembled from all discovered states
///
/// # Complexity
///
/// - **Time:** O(2^V / P) worst case, O((V + E) / P) typical case
/// - **Space:** O(2^V) for subset storage (same as sequential)
///
/// # Examples
///
/// ```rust
/// use arcweight::prelude::*;
/// use arcweight::algorithms::parallel::determinize_parallel;
///
/// 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());
/// // Non-deterministic: two arcs with same input label
/// fst.add_arc(s0, Arc::new(1, 1, TropicalWeight::new(1.0), s1));
/// fst.add_arc(s0, Arc::new(1, 1, TropicalWeight::new(2.0), s2));
/// fst.add_arc(s1, Arc::new(2, 2, TropicalWeight::new(1.0), s2));
///
/// let det: VectorFst<TropicalWeight> = determinize_parallel(&fst).unwrap();
/// // Result is deterministic
/// ```
#[cfg(feature = "parallel")]
pub fn determinize_parallel<W, F, M>(fst: &F) -> Result<M>
where
    W: crate::semiring::DivisibleSemiring + std::hash::Hash + Eq + Ord + Send + Sync,
    F: Fst<W> + Sync,
    M: MutableFst<W> + Default + Send,
{
    use std::collections::BTreeMap;
    use std::sync::atomic::{AtomicU32, Ordering};

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

    // Weighted subset type
    #[derive(Clone, Debug, PartialEq, Eq, Hash)]
    struct WeightedSubset<W: Semiring> {
        states: BTreeMap<StateId, W>,
    }

    impl<W: Semiring> WeightedSubset<W> {
        fn new() -> Self {
            Self {
                states: BTreeMap::new(),
            }
        }

        fn insert(&mut self, state: StateId, weight: W) {
            self.states
                .entry(state)
                .and_modify(|w| w.plus_assign(&weight))
                .or_insert(weight);
        }

        fn normalize(&mut self) -> Option<W>
        where
            W: crate::semiring::DivisibleSemiring + Ord,
        {
            if self.states.is_empty() {
                return None;
            }

            let min_weight = self.states.values().min()?.clone();

            if <W as num_traits::Zero>::is_zero(&min_weight) {
                return None;
            }

            for weight in self.states.values_mut() {
                match weight.divide(&min_weight) {
                    Some(normalized) => *weight = normalized,
                    None => return None,
                }
            }

            Some(min_weight)
        }
    }

    // Thread-safe data structures
    let subset_map: StdArc<RwLock<FxHashMap<WeightedSubset<W>, StateId>>> =
        StdArc::new(RwLock::new(FxHashMap::default()));
    let next_state_id = AtomicU32::new(0);

    // Work queue: subsets to process
    let work_queue: StdArc<Mutex<Vec<(WeightedSubset<W>, StateId)>>> =
        StdArc::new(Mutex::new(Vec::new()));

    // Collected results
    let arc_data: StdArc<Mutex<Vec<(StateId, crate::fst::Label, W, StateId)>>> =
        StdArc::new(Mutex::new(Vec::new()));
    let final_data: StdArc<Mutex<Vec<(StateId, W)>>> = StdArc::new(Mutex::new(Vec::new()));

    // Create initial subset
    let mut start_subset = WeightedSubset::new();
    start_subset.insert(start, W::one());

    // Register start state
    let start_new = next_state_id.fetch_add(1, Ordering::Relaxed);
    {
        let mut map = subset_map.write().unwrap();
        map.insert(start_subset.clone(), start_new);
    }

    // Add to work queue
    work_queue.lock().unwrap().push((start_subset, start_new));

    // Parallel processing with work stealing
    let num_threads = rayon::current_num_threads();
    let processed_count = AtomicU32::new(0);
    let total_work_estimate = AtomicU32::new(1); // At least start subset

    rayon::scope(|s| {
        for _ in 0..num_threads {
            let subset_map = StdArc::clone(&subset_map);
            let work_queue = StdArc::clone(&work_queue);
            let arc_data = StdArc::clone(&arc_data);
            let final_data = StdArc::clone(&final_data);
            let next_state_id = &next_state_id;
            let processed_count = &processed_count;
            let total_work_estimate = &total_work_estimate;

            s.spawn(move |_| {
                loop {
                    // Try to steal work
                    let work_item = {
                        let mut queue = work_queue.lock().unwrap();
                        queue.pop()
                    };

                    match work_item {
                        Some((subset, current_state)) => {
                            // Process this subset
                            let mut transitions: FxHashMap<crate::fst::Label, WeightedSubset<W>> =
                                FxHashMap::default();
                            let mut final_weight = W::zero();

                            for (&state, weight) in &subset.states {
                                // Accumulate final weights
                                if let Some(fw) = fst.final_weight(state) {
                                    final_weight.plus_assign(&weight.times(fw));
                                }

                                // Process arcs
                                for arc in fst.arcs(state) {
                                    let next_weight = weight.times(&arc.weight);
                                    transitions
                                        .entry(arc.ilabel)
                                        .or_insert_with(WeightedSubset::new)
                                        .insert(arc.nextstate, next_weight);
                                }
                            }

                            // Record final weight
                            if !<W as num_traits::Zero>::is_zero(&final_weight) {
                                final_data
                                    .lock()
                                    .unwrap()
                                    .push((current_state, final_weight));
                            }

                            // Process transitions
                            for (label, mut next_subset) in transitions {
                                if let Some(norm_weight) = next_subset.normalize() {
                                    // Check if subset exists
                                    let existing = {
                                        let map = subset_map.read().unwrap();
                                        map.get(&next_subset).copied()
                                    };

                                    let next_state = match existing {
                                        Some(state) => state,
                                        None => {
                                            // Create new state
                                            let mut map = subset_map.write().unwrap();
                                            // Double-check after acquiring write lock
                                            if let Some(&state) = map.get(&next_subset) {
                                                state
                                            } else {
                                                let new_state =
                                                    next_state_id.fetch_add(1, Ordering::Relaxed);
                                                map.insert(next_subset.clone(), new_state);
                                                // Add to work queue
                                                work_queue
                                                    .lock()
                                                    .unwrap()
                                                    .push((next_subset, new_state));
                                                // Use Release to ensure work item is visible before count update
                                                total_work_estimate.fetch_add(1, Ordering::Release);
                                                new_state
                                            }
                                        }
                                    };

                                    // Record arc
                                    arc_data.lock().unwrap().push((
                                        current_state,
                                        label,
                                        norm_weight,
                                        next_state,
                                    ));
                                }
                            }

                            // Use Release ordering to ensure all state updates are visible
                            processed_count.fetch_add(1, Ordering::Release);
                        }
                        None => {
                            // No work available, check if we're done
                            // Use Acquire ordering to see all updates from other threads
                            let processed = processed_count.load(Ordering::Acquire);
                            let total = total_work_estimate.load(Ordering::Acquire);

                            if processed >= total {
                                // Double-check queue is empty while holding the lock
                                // This ensures we don't miss work items added concurrently
                                let queue_empty = work_queue.lock().unwrap().is_empty();
                                if queue_empty {
                                    // Final check: re-read counters with lock held
                                    let final_processed = processed_count.load(Ordering::Acquire);
                                    let final_total = total_work_estimate.load(Ordering::Acquire);
                                    if final_processed >= final_total {
                                        break;
                                    }
                                }
                            }

                            // Brief yield to avoid busy-waiting
                            std::thread::yield_now();
                        }
                    }
                }
            });
        }
    });

    // Build result FST
    let num_states = next_state_id.load(Ordering::Relaxed) as usize;
    let mut result = M::default();

    for _ in 0..num_states {
        result.add_state();
    }

    result.set_start(start_new);

    // Add final weights
    for (state, weight) in final_data.lock().unwrap().drain(..) {
        result.set_final(state, weight);
    }

    // Add arcs
    for (source, label, weight, target) in arc_data.lock().unwrap().drain(..) {
        result.add_arc(source, Arc::new(label, label, weight, target));
    }

    Ok(result)
}

/// Non-parallel fallback for determinize_parallel
#[cfg(not(feature = "parallel"))]
pub fn determinize_parallel<W, F, M>(fst: &F) -> Result<M>
where
    W: crate::semiring::DivisibleSemiring + std::hash::Hash + Eq + Ord,
    F: Fst<W>,
    M: MutableFst<W> + Default,
{
    crate::algorithms::determinize(fst)
}

/// Parallel minimization using partition refinement
///
/// This algorithm parallelizes Hopcroft's partition refinement algorithm
/// by processing refinement blocks in parallel.
///
/// # Algorithm
///
/// 1. Initial partition: {final states} ∪ {non-final states}
/// 2. Worklist contains blocks to be processed
/// 3. Workers process blocks in parallel:
///    - For each block B, split other blocks based on transitions into B
///    - Add resulting smaller blocks to worklist
/// 4. Build minimal FST from final partition
///
/// # Complexity
///
/// - **Time:** O(n log n / P) where n = number of states, P = processors
/// - **Space:** O(n) for partition storage
#[cfg(feature = "parallel")]
pub fn minimize_parallel<W, F, M>(fst: &F) -> Result<M>
where
    W: Semiring + std::hash::Hash + Eq + Send + Sync,
    F: Fst<W> + Sync,
    M: MutableFst<W> + Default + Send,
{
    use std::sync::atomic::{AtomicUsize, Ordering};

    let num_states = fst.num_states();
    if num_states == 0 {
        return Ok(M::default());
    }

    // Initial partition: separate final and non-final states
    let mut final_states: FxHashSet<StateId> = FxHashSet::default();
    let mut non_final_states: FxHashSet<StateId> = FxHashSet::default();

    // Collect final state weights for grouping
    let mut final_weight_groups: FxHashMap<u64, Vec<StateId>> = FxHashMap::default();

    for state in fst.states() {
        if let Some(weight) = fst.final_weight(state) {
            final_states.insert(state);
            // Hash the weight for grouping (states with same final weight)
            let weight_hash = {
                use std::hash::Hasher;
                let mut hasher = rustc_hash::FxHasher::default();
                std::hash::Hash::hash(weight, &mut hasher);
                hasher.finish()
            };
            final_weight_groups
                .entry(weight_hash)
                .or_default()
                .push(state);
        } else {
            non_final_states.insert(state);
        }
    }

    // Build initial partition
    let mut partition: Vec<FxHashSet<StateId>> = Vec::new();

    // Add final state groups
    for (_hash, states) in final_weight_groups {
        if !states.is_empty() {
            partition.push(states.into_iter().collect());
        }
    }

    // Add non-final states
    if !non_final_states.is_empty() {
        partition.push(non_final_states);
    }

    // Build reverse adjacency: for each state, list of (source_state, label) pairs
    let reverse_adj: Vec<Vec<(StateId, crate::fst::Label)>> = {
        let mut rev = vec![Vec::new(); num_states];
        for state in fst.states() {
            for arc in fst.arcs(state) {
                rev[arc.nextstate as usize].push((state, arc.ilabel));
            }
        }
        rev
    };

    // State to partition index mapping
    let state_to_partition: StdArc<RwLock<Vec<usize>>> =
        StdArc::new(RwLock::new(vec![0; num_states]));

    // Update initial mapping
    {
        let mut map = state_to_partition.write().unwrap();
        for (idx, block) in partition.iter().enumerate() {
            for &state in block {
                map[state as usize] = idx;
            }
        }
    }

    // Worklist of block indices to process
    let worklist: StdArc<Mutex<Vec<usize>>> =
        StdArc::new(Mutex::new((0..partition.len()).collect()));

    // Partition storage with synchronization
    let partition_storage: StdArc<RwLock<Vec<FxHashSet<StateId>>>> =
        StdArc::new(RwLock::new(partition));

    let refinement_done = std::sync::atomic::AtomicBool::new(false);
    let active_workers = AtomicUsize::new(0);

    // Parallel partition refinement
    rayon::scope(|s| {
        let num_threads = rayon::current_num_threads();

        for _ in 0..num_threads {
            let state_to_partition = StdArc::clone(&state_to_partition);
            let worklist = StdArc::clone(&worklist);
            let partition_storage = StdArc::clone(&partition_storage);
            let reverse_adj = &reverse_adj;
            let refinement_done = &refinement_done;
            let active_workers = &active_workers;

            s.spawn(move |_| {
                loop {
                    // Try to get work
                    let block_idx = {
                        let mut wl = worklist.lock().unwrap();
                        wl.pop()
                    };

                    match block_idx {
                        Some(idx) => {
                            active_workers.fetch_add(1, Ordering::Relaxed);

                            // Get the block to process
                            let block: Vec<StateId> = {
                                let parts = partition_storage.read().unwrap();
                                if idx < parts.len() {
                                    parts[idx].iter().copied().collect()
                                } else {
                                    Vec::new()
                                }
                            };

                            if block.is_empty() {
                                active_workers.fetch_sub(1, Ordering::Relaxed);
                                continue;
                            }

                            // Find all states that can reach this block, grouped by label
                            let mut predecessors_by_label: FxHashMap<
                                crate::fst::Label,
                                FxHashSet<StateId>,
                            > = FxHashMap::default();

                            for &target_state in &block {
                                for &(source_state, label) in &reverse_adj[target_state as usize] {
                                    predecessors_by_label
                                        .entry(label)
                                        .or_default()
                                        .insert(source_state);
                                }
                            }

                            // For each label, refine other blocks
                            for (_label, predecessors) in predecessors_by_label {
                                // Find blocks to split
                                let blocks_to_check: Vec<usize> = {
                                    let map = state_to_partition.read().unwrap();
                                    let mut block_indices: FxHashSet<usize> = FxHashSet::default();
                                    for &pred in &predecessors {
                                        block_indices.insert(map[pred as usize]);
                                    }
                                    block_indices.into_iter().collect()
                                };

                                for check_idx in blocks_to_check {
                                    // Get states in this block
                                    let block_states: Vec<StateId> = {
                                        let parts = partition_storage.read().unwrap();
                                        if check_idx < parts.len() {
                                            parts[check_idx].iter().copied().collect()
                                        } else {
                                            continue;
                                        }
                                    };

                                    // Split: states that can reach target block vs those that can't
                                    let can_reach: FxHashSet<StateId> = block_states
                                        .iter()
                                        .filter(|s| predecessors.contains(s))
                                        .copied()
                                        .collect();

                                    let cannot_reach: FxHashSet<StateId> = block_states
                                        .iter()
                                        .filter(|s| !predecessors.contains(s))
                                        .copied()
                                        .collect();

                                    // If both sets are non-empty, we have a split
                                    if !can_reach.is_empty() && !cannot_reach.is_empty() {
                                        // Lock ordering: always acquire partition_storage before state_to_partition
                                        // to prevent deadlocks when multiple threads need both locks.
                                        let mut parts = partition_storage.write().unwrap();
                                        let mut map = state_to_partition.write().unwrap();

                                        // Keep smaller set in original block, create new block for larger
                                        let (keep, split) = if can_reach.len() <= cannot_reach.len()
                                        {
                                            (can_reach, cannot_reach)
                                        } else {
                                            (cannot_reach, can_reach)
                                        };

                                        // Update original block
                                        if check_idx < parts.len() {
                                            parts[check_idx] = keep;
                                        }

                                        // Create new block
                                        let new_idx = parts.len();
                                        parts.push(split.clone());

                                        // Update state mapping
                                        for &state in &split {
                                            map[state as usize] = new_idx;
                                        }

                                        // Add new block to worklist
                                        worklist.lock().unwrap().push(new_idx);
                                    }
                                }
                            }

                            active_workers.fetch_sub(1, Ordering::Relaxed);
                        }
                        None => {
                            // No work available
                            if active_workers.load(Ordering::Relaxed) == 0 {
                                let wl = worklist.lock().unwrap();
                                if wl.is_empty() {
                                    refinement_done.store(true, Ordering::Relaxed);
                                    break;
                                }
                            }

                            if refinement_done.load(Ordering::Relaxed) {
                                break;
                            }

                            std::thread::yield_now();
                        }
                    }
                }
            });
        }
    });

    // Build minimal FST from final partition
    let final_partition = partition_storage.read().unwrap();
    let state_map = state_to_partition.read().unwrap();

    let mut result = M::default();

    // Create one state per non-empty partition block
    let mut block_to_new_state: Vec<Option<StateId>> = vec![None; final_partition.len()];
    let mut new_state_count = 0;

    for (idx, block) in final_partition.iter().enumerate() {
        if !block.is_empty() {
            block_to_new_state[idx] = Some(result.add_state());
            new_state_count += 1;
        }
    }

    if new_state_count == 0 {
        return Ok(result);
    }

    // Set start state
    if let Some(start) = fst.start() {
        let start_block = state_map[start as usize];
        if let Some(new_start) = block_to_new_state[start_block] {
            result.set_start(new_start);
        }
    }

    // Add arcs and final weights (use representative state from each block)
    for (idx, block) in final_partition.iter().enumerate() {
        if block.is_empty() {
            continue;
        }

        let new_state = match block_to_new_state[idx] {
            Some(s) => s,
            None => continue,
        };

        // Use first state as representative
        let representative = *block.iter().next().unwrap();

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

        // Add arcs (deduplicated)
        let mut seen_arcs: FxHashSet<(crate::fst::Label, crate::fst::Label, StateId)> =
            FxHashSet::default();

        for arc in fst.arcs(representative) {
            let target_block = state_map[arc.nextstate as usize];
            if let Some(new_target) = block_to_new_state[target_block] {
                let arc_key = (arc.ilabel, arc.olabel, new_target);
                if seen_arcs.insert(arc_key) {
                    result.add_arc(
                        new_state,
                        Arc::new(arc.ilabel, arc.olabel, arc.weight.clone(), new_target),
                    );
                }
            }
        }
    }

    Ok(result)
}

/// Non-parallel fallback for minimize_parallel
#[cfg(not(feature = "parallel"))]
pub fn minimize_parallel<W, F, M>(fst: &F) -> Result<M>
where
    W: crate::semiring::DivisibleSemiring + std::hash::Hash + Eq + Ord,
    F: Fst<W>,
    M: MutableFst<W> + Default,
{
    crate::algorithms::minimize(fst)
}

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

    #[test]
    fn test_map_weights_parallel() {
        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(1.0));
        fst.add_arc(s0, Arc::new(1, 1, TropicalWeight::new(2.0), s1));

        let result: VectorFst<TropicalWeight> =
            map_weights_parallel(&fst, |w| TropicalWeight::new(w.value() * 2.0)).unwrap();

        assert_eq!(result.num_states(), 2);
        let arcs: Vec<_> = result.arcs(s0).collect();
        assert_eq!(arcs[0].weight, TropicalWeight::new(4.0));
        assert_eq!(result.final_weight(s1), Some(&TropicalWeight::new(2.0)));
    }

    #[test]
    fn test_parallel_state_map() {
        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.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(3.0), s2));

        let arc_counts: Vec<usize> = parallel_state_map(&fst, |state, f| f.num_arcs(state));

        assert_eq!(arc_counts, vec![2, 1, 0]);
    }

    #[test]
    fn test_collect_arcs_parallel() {
        let mut fst = VectorFst::<TropicalWeight>::new();
        let s0 = fst.add_state();
        let s1 = fst.add_state();
        fst.set_start(s0);
        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), s1));

        let arcs = collect_arcs_parallel(&fst);

        assert_eq!(arcs.len(), 2);
        assert_eq!(arcs[0].1.len(), 2);
        assert_eq!(arcs[1].1.len(), 0);
    }

    #[test]
    fn test_compose_parallel_basic() {
        // First FST: input -> intermediate
        let mut fst1 = VectorFst::<TropicalWeight>::new();
        let s0 = fst1.add_state();
        let s1 = fst1.add_state();
        fst1.set_start(s0);
        fst1.set_final(s1, TropicalWeight::one());
        fst1.add_arc(s0, Arc::new(1, 2, TropicalWeight::new(1.0), s1));

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

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

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

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

    #[test]
    fn test_compose_parallel_multiple_paths() {
        let mut fst1 = VectorFst::<TropicalWeight>::new();
        let s0 = fst1.add_state();
        let s1 = fst1.add_state();
        fst1.set_start(s0);
        fst1.set_final(s1, TropicalWeight::one());
        fst1.add_arc(s0, Arc::new(1, 2, TropicalWeight::new(1.0), s1));
        fst1.add_arc(s0, Arc::new(1, 3, TropicalWeight::new(1.5), s1));

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

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

        // Should have both paths
        let mut found_path1 = false;
        let mut found_path2 = false;
        for state in composed.states() {
            for arc in composed.arcs(state) {
                if arc.ilabel == 1 && arc.olabel == 4 {
                    found_path1 = true;
                }
                if arc.ilabel == 1 && arc.olabel == 5 {
                    found_path2 = true;
                }
            }
        }
        assert!(found_path1, "Should find first path");
        assert!(found_path2, "Should find second path");
    }

    #[test]
    fn test_compose_parallel_with_epsilon() {
        let mut fst1 = VectorFst::<TropicalWeight>::new();
        let s0 = fst1.add_state();
        let s1 = fst1.add_state();
        let s2 = fst1.add_state();
        fst1.set_start(s0);
        fst1.set_final(s2, TropicalWeight::one());
        fst1.add_arc(s0, Arc::epsilon(TropicalWeight::new(0.5), s1));
        fst1.add_arc(s1, Arc::new(1, 2, TropicalWeight::new(1.0), s2));

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

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

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

    #[test]
    fn test_compose_parallel_no_match() {
        let mut fst1 = VectorFst::<TropicalWeight>::new();
        let s0 = fst1.add_state();
        let s1 = fst1.add_state();
        fst1.set_start(s0);
        fst1.set_final(s1, TropicalWeight::one());
        fst1.add_arc(s0, Arc::new(1, 2, TropicalWeight::new(1.0), s1));

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

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

        // Should have no arcs (no matching labels)
        let total_arcs: usize = composed.states().map(|s| composed.num_arcs(s)).sum();
        assert_eq!(total_arcs, 0);
    }

    #[test]
    fn test_shortest_distance_parallel_basic() {
        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 distances = shortest_distance_parallel(&fst).unwrap();

        assert_eq!(distances.len(), 3);
        // Start state should have distance 0 (one in tropical)
        assert_eq!(*distances[0].value(), 0.0);
        // Distance to s1: 1.0
        assert_eq!(*distances[1].value(), 1.0);
        // Distance to s2: 1.0 + 2.0 = 3.0
        assert_eq!(*distances[2].value(), 3.0);
    }

    #[test]
    fn test_shortest_distance_parallel_multiple_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(s2, TropicalWeight::one());
        // Two paths to s2: direct (5.0) and via s1 (1.0 + 2.0 = 3.0)
        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));
        fst.add_arc(s0, Arc::new(3, 3, TropicalWeight::new(5.0), s2));

        let distances = shortest_distance_parallel(&fst).unwrap();

        // Shortest distance to s2 should be min(5.0, 3.0) = 3.0
        assert_eq!(*distances[2].value(), 3.0);
    }

    #[test]
    fn test_delta_stepping_basic() {
        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 distances = delta_stepping_shortest_distance(&fst, 1.0).unwrap();

        assert_eq!(distances.len(), 3);
        // Start state should have distance 0
        assert_eq!(*distances[0].value(), 0.0);
        // Distance to s1: 1.0
        assert_eq!(*distances[1].value(), 1.0);
        // Distance to s2: 3.0
        assert_eq!(*distances[2].value(), 3.0);
    }

    #[test]
    fn test_delta_stepping_multiple_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(s2, TropicalWeight::one());
        // Two paths: direct (5.0) and via s1 (1.0 + 2.0 = 3.0)
        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));
        fst.add_arc(s0, Arc::new(3, 3, TropicalWeight::new(5.0), s2));

        let distances = delta_stepping_shortest_distance(&fst, 2.0).unwrap();

        // Should find minimum: 3.0
        assert_eq!(*distances[2].value(), 3.0);
    }

    #[test]
    fn test_delta_stepping_empty_fst() {
        let fst = VectorFst::<TropicalWeight>::new();
        // Empty FST: may return error or empty result depending on implementation
        let result = delta_stepping_shortest_distance(&fst, 1.0);
        if let Ok(distances) = result {
            assert!(distances.is_empty());
        }
        // Err case is also acceptable - no start state
    }

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

        // Non-deterministic: two arcs with same input label
        fst.add_arc(s0, Arc::new(1, 1, TropicalWeight::new(1.0), s1));
        fst.add_arc(s0, Arc::new(1, 1, TropicalWeight::new(2.0), s2));
        fst.add_arc(s1, Arc::new(2, 2, TropicalWeight::new(1.0), s2));

        let det: VectorFst<TropicalWeight> = determinize_parallel(&fst).unwrap();

        // Check determinism: no state should have multiple arcs with same input label
        for state in det.states() {
            let mut seen_labels = std::collections::HashSet::new();
            for arc in det.arcs(state) {
                assert!(
                    seen_labels.insert(arc.ilabel),
                    "Found duplicate input label {} from state {}",
                    arc.ilabel,
                    state
                );
            }
        }

        // Should preserve language
        assert!(det.start().is_some());
        assert!(det.num_states() > 0);
    }

    #[test]
    fn test_determinize_parallel_already_deterministic() {
        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(1.0), s1));

        let det: VectorFst<TropicalWeight> = determinize_parallel(&fst).unwrap();

        // Should be similar to original
        assert_eq!(det.num_states(), fst.num_states());
        assert!(det.start().is_some());
    }

    #[test]
    fn test_minimize_parallel_simple() {
        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::one());
        fst.set_final(s3, TropicalWeight::one());

        // Create redundant paths
        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(1, 1, TropicalWeight::new(1.0), s3));

        let minimized: VectorFst<TropicalWeight> = minimize_parallel(&fst).unwrap();

        // Result should be valid (Brzozowski's algorithm produces minimal but may differ from input)
        assert!(minimized.start().is_some());
        assert!(minimized.num_states() > 0);
    }

    #[test]
    fn test_minimize_parallel_already_minimal() {
        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(1.0), s1));

        let minimized: VectorFst<TropicalWeight> = minimize_parallel(&fst).unwrap();

        // Result should be valid (Brzozowski's algorithm may produce different structures)
        assert!(minimized.start().is_some());
        assert!(minimized.num_states() > 0);
    }

    #[test]
    fn test_minimize_parallel_empty() {
        let fst = VectorFst::<TropicalWeight>::new();
        let minimized: VectorFst<TropicalWeight> = minimize_parallel(&fst).unwrap();
        assert_eq!(minimized.num_states(), 0);
    }
}