arcweight 0.3.0

//! Weighted determinization via subset construction.
//!
//! This module implements determinization of weighted finite-state transducers using
//! the weighted subset construction algorithm. Determinization ensures each state has
//! at most one outgoing arc per input label, enabling efficient $`O(|x|)`$ lookup.
//!
//! # Mathematical Definition
//!
//! A weighted automaton $`A = (Q, \Sigma, E, i, F, \lambda, \rho)`$ is deterministic if:
//! - There is exactly one initial state
//! - For each state $`q`$ and label $`a`$, at most one transition $`(q, a, w, q')`$ exists
//!
//! Determinization constructs an equivalent deterministic automaton by tracking
//! subsets of states with accumulated weights.
//!
//! # Complexity
//!
//! | Case | Time | Space |
//! |------|------|-------|
//! | Worst | $`O(2^V)`$ | $`O(2^V)`$ |
//! | Typical | $`O(V + E)`$ | $`O(V + E)`$ |
//!
//! The exponential worst case is rare in practice; most FSTs from speech/NLP
//! applications have polynomial blowup.
//!
//! # Semiring Requirements
//!
//! Requires [`DivisibleSemiring`] for weight normalization during subset construction:
//!
//! | Semiring | Supported | Notes |
//! |----------|-----------|-------|
//! | [`TropicalWeight`] | Yes | Natural division |
//! | [`LogWeight`] | Yes | Division in log space |
//! | [`ProbabilityWeight`] | No | Not weakly left divisible |
//! | [`BooleanWeight`] | No | No natural division |
//!
//! # Algorithm
//!
//! Weighted subset construction (Mohri, 1997):
//!
//! 1. **Initialize:** Create start subset $`\{(q_0, \bar{1})\}`$
//! 2. **Expand:** For each subset $`S`$:
//!    - Group arcs by $`(\text{ilabel}, \text{olabel})`$
//!    - For each label pair, compute destination subset
//!    - **Normalize:** Divide all weights by minimum to bound growth
//!    - Create arc with normalization weight
//! 3. **Terminate:** When all subsets processed
//!
//! # Transducer Support
//!
//! Unlike acceptor determinization, this implementation groups by
//! $`(\text{ilabel}, \text{olabel})`$ pairs, preserving output labels correctly.
//!
//! # Example
//!
//! ```rust
//! use arcweight::prelude::*;
//!
//! // Nondeterministic FST
//! let mut nfst = VectorFst::<TropicalWeight>::new();
//! let s0 = nfst.add_state();
//! let s1 = nfst.add_state();
//! let s2 = nfst.add_state();
//! nfst.set_start(s0);
//! nfst.set_final(s1, TropicalWeight::one());
//! nfst.set_final(s2, TropicalWeight::one());
//!
//! // Two arcs with same input label (nondeterminism)
//! nfst.add_arc(s0, Arc::new(1, 1, TropicalWeight::new(0.5), s1));
//! nfst.add_arc(s0, Arc::new(1, 1, TropicalWeight::new(0.7), s2));
//!
//! // Determinize
//! let dfst: VectorFst<TropicalWeight> = determinize(&nfst)?;
//!
//! // Result has one arc per input label from each state
//! # Ok::<(), arcweight::Error>(())
//! ```
//!
//! # References
//!
//! \[1\] Mohri, M. 1997. Finite-state transducers in language and speech processing.
//!     *Computational Linguistics* 23, 2 (June 1997), 269-311.
//!
//! \[2\] Mohri, M. 2009. Weighted automata algorithms. In *Handbook of Weighted
//!     Automata*, M. Droste, W. Kuich, and H. Vogler, Eds. Springer, 213-254.
//!     <https://doi.org/10.1007/978-3-642-01492-5_6>
//!
//! \[3\] Allauzen, C. and Mohri, M. 2003. Efficient algorithms for testing the twins
//!     property. *Journal of Automata, Languages and Combinatorics* 8, 2, 117-144.
//!
//! [`DivisibleSemiring`]: crate::semiring::DivisibleSemiring
//! [`TropicalWeight`]: crate::semiring::TropicalWeight
//! [`LogWeight`]: crate::semiring::LogWeight
//! [`ProbabilityWeight`]: crate::semiring::ProbabilityWeight
//! [`BooleanWeight`]: crate::semiring::BooleanWeight

use crate::arc::Arc;
use crate::fst::{Fst, Label, MutableFst, StateId};
use crate::semiring::{DivisibleSemiring, Semiring};
use crate::{Error, Result};
use core::hash::Hash;
use rustc_hash::FxHashMap;
use std::collections::BTreeMap;

/// Weighted subset (determinization state)
#[derive(Clone, Debug, PartialEq)]
struct WeightedSubset<W: Semiring> {
    /// States with their weights
    states: BTreeMap<StateId, W>,
}

impl<W: Semiring> Eq for WeightedSubset<W> where W: Eq {}

impl<W: Semiring> std::hash::Hash for WeightedSubset<W>
where
    W: std::hash::Hash,
{
    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
        self.states.hash(state);
    }
}

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

    /// Normalize weights by dividing by the minimum weight
    ///
    /// This prevents exponential weight growth during subset construction.
    /// Returns the normalization factor (minimum weight) if successful.
    ///
    /// # Edge Cases
    ///
    /// - If all weights are infinity (zero element), returns `None` because
    ///   division by infinity is undefined. This indicates that there are
    ///   no valid paths from this subset, so the transition should be skipped.
    /// - If the subset is empty, returns `None`.
    ///
    /// # Examples
    ///
    /// For tropical semiring:
    /// - If weights are [5.0, 3.0, 7.0], min is 3.0
    /// - After normalization: [2.0, 0.0, 4.0] (5-3, 3-3, 7-3)
    /// - Returns Some(3.0)
    ///
    /// If all weights are infinity:
    /// - Returns None (no valid path)
    fn normalize(&mut self) -> Option<W>
    where
        W: DivisibleSemiring + Ord,
    {
        // Check if subset is empty
        if self.states.is_empty() {
            return None;
        }

        // Find minimum weight
        let min_weight = self.states.values().min()?.clone();

        // Check if minimum weight is the semiring's additive identity (zero element).
        // In semiring terminology:
        //   - Zero element = additive identity (a + 0 = a) = "worst" or "no path"
        //   - One element = multiplicative identity (a * 1 = a) = "neutral cost"
        //
        // For tropical semiring: zero = f32::INFINITY (worst cost), one = 0.0 (no cost)
        // For log semiring: zero = f32::INFINITY (log(0)), one = 0.0 (log(1))
        //
        // The `is_zero` check tests for the semiring zero, NOT the floating-point value 0.0.
        // Division by the zero element is undefined (no valid normalization possible).
        if <W as num_traits::Zero>::is_zero(&min_weight) {
            return None;
        }

        // Divide all weights by minimum
        for weight in self.states.values_mut() {
            match weight.divide(&min_weight) {
                Some(normalized) => *weight = normalized,
                None => {
                    // Division failed (shouldn't happen if min_weight is not zero)
                    // But handle gracefully by returning None
                    return None;
                }
            }
        }

        Some(min_weight)
    }
}

/// Determinize a weighted FST using subset construction
///
/// Converts a nondeterministic FST into a deterministic one that accepts the same
/// weighted language. Ensures each state has at most one outgoing arc per input label,
/// combining paths via semiring addition (⊕) and preserving weights through normalization.
///
/// Requires [`DivisibleSemiring`] for weight normalization during subset construction.
/// Compatible with [`TropicalWeight`] and [`LogWeight`].
///
/// # Complexity
///
/// - **Time:** O(2^V) worst case, O(V + E) typical case
///   - V = number of states in input FST
///   - E = number of arcs in input FST
///   - Worst case: exponential blowup (rare, requires extensive nondeterminism)
///   - Typical case: linear or near-linear with sparse nondeterminism
/// - **Space:** O(2^V) for subset storage (FxHashMap of subsets)
///   - Each unique subset becomes one state in result
///   - Memory dominated by subset representation
///
/// # Algorithm
///
/// Weighted subset construction with normalization (Mohri, 1997):
/// 1. Initialize: Create start state from {start_state: 1̄}
/// 2. Process queue of (subset, state) pairs:
///    - Compute final weight: w_f = ⊕_{s ∈ subset} w_s ⊗ final_weight(s)
///    - Group arcs by input label: for each label ℓ
///      * Compute destination subset: {(next_s, w_s ⊗ arc.weight) | arc with label ℓ}
///      * Normalize: divide all weights by min weight in subset
///      * Create/lookup state for normalized subset
///      * Add arc with label ℓ and normalization weight
/// 3. Return deterministic FST with same weighted language
///
/// # Performance Notes
///
/// - **Subset explosion:** Rare but possible with extensive nondeterminism
/// - **Early stopping:** Consider weight pruning for very large FSTs
/// - **Normalization overhead:** Division per subset, critical for correctness
/// - **Label sparsity:** Fewer labels per state improves performance
/// - **Preprocessing:** Remove epsilon arcs before determinization for best results
/// - **Memory pattern:** FxHashMap lookups dominate; good cache locality with BTreeMap subsets
///
/// [`DivisibleSemiring`]: crate::semiring::DivisibleSemiring
/// [`TropicalWeight`]: crate::semiring::TropicalWeight
/// [`LogWeight`]: crate::semiring::LogWeight
///
/// # Examples
///
/// ## Basic Determinization
///
/// ```rust
/// use arcweight::prelude::*;
///
/// // Create nondeterministic FST with ambiguous 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::new(0.5));
/// fst.set_final(s2, TropicalWeight::new(0.3));
///
/// // Two arcs with same input label - creates nondeterminism
/// fst.add_arc(s0, Arc::new('a' as u32, 'x' as u32, TropicalWeight::new(0.2), s1));
/// fst.add_arc(s0, Arc::new('a' as u32, 'y' as u32, TropicalWeight::new(0.4), s2));
///
/// // Determinize resolves nondeterminism
/// let det_fst: VectorFst<TropicalWeight> = determinize(&fst).unwrap();
///
/// // Determinization may create more or fewer states
/// assert!(det_fst.num_states() > 0);
/// ```
///
/// ## Real-World Application: Speech Recognition
///
/// ```rust
/// use arcweight::prelude::*;
///
/// // Determinize pronunciation dictionary for efficient lookup
/// fn optimize_pronunciation_dict(
///     dict: &VectorFst<TropicalWeight>
/// ) -> std::result::Result<VectorFst<TropicalWeight>, Box<dyn std::error::Error>> {
///     
///     // Step 1: Determinize to resolve pronunciation ambiguities
///     let det_dict: VectorFst<TropicalWeight> = determinize(dict)?;
///     
///     // Step 2: Minimize to reduce memory footprint
///     let min_dict: VectorFst<TropicalWeight> = minimize(&det_dict)?;
///     
///     println!("Original: {} states, Optimized: {} states",
///              dict.num_states(), min_dict.num_states());
///     
///     Ok(min_dict)
/// }
/// ```
///
/// ## Checking Determinism
///
/// ```rust
/// use arcweight::prelude::*;
///
/// fn ensure_deterministic<W: DivisibleSemiring + Ord + std::hash::Hash + Eq>(
///     fst: &VectorFst<W>
/// ) -> std::result::Result<VectorFst<W>, Box<dyn std::error::Error>> {
///     
///     // Simple example: always determinize for demonstration
///     println!("Determinizing FST...");
///     let det_fst: VectorFst<W> = determinize(fst)?;
///     println!("Original: {} states, Determinized: {} states",
///              fst.num_states(), det_fst.num_states());
///     
///     Ok(det_fst)
/// }
/// ```
///
/// # Errors
///
/// Returns [`Error::Algorithm`] if:
/// - The input FST has no start state or is malformed
/// - Memory allocation fails during subset construction
/// - The semiring doesn't support required division operations
/// - Weight normalization fails due to division by zero
/// - Subset construction creates invalid state combinations
///
/// # References
///
/// \[1\] Mohri, M. 1997. Finite-state transducers in language and speech processing.
///     *Computational Linguistics* 23, 2 (June 1997), 269-311.
///
/// \[2\] Mohri, M. 2009. Weighted automata algorithms. In *Handbook of Weighted
///     Automata*, Springer, 213-254. <https://doi.org/10.1007/978-3-642-01492-5_6>
///
/// # See Also
///
/// - [`minimize`] - Reduce deterministic FST size (often applied after determinization)
/// - [`DivisibleSemiring`] - Required trait for weight normalization
/// - [`TropicalWeight`] - Compatible semiring (min-plus algebra)
/// - [`LogWeight`] - Compatible semiring (log-space probabilities)
/// - [`compose`] - Often combined with determinization in FST pipelines
/// - [Semiring trait](crate::semiring::Semiring) - Mathematical foundation
///
/// [`minimize`]: crate::algorithms::minimize::minimize
/// [`compose`]: crate::algorithms::compose::compose
pub fn determinize<W, F, M>(fst: &F) -> Result<M>
where
    W: DivisibleSemiring + Hash + Eq + Ord,
    F: Fst<W>,
    M: MutableFst<W> + Default,
{
    let start = fst
        .start()
        .ok_or_else(|| Error::Algorithm("FST has no start state".into()))?;

    let mut result = M::default();
    let mut subset_map = FxHashMap::default();
    let mut queue = Vec::new();

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

    // create start state
    let start_new = result.add_state();
    result.set_start(start_new);
    subset_map.insert(start_subset.clone(), start_new);
    queue.push((start_subset, start_new));

    // process subsets
    while let Some((subset, current_state)) = queue.pop() {
        // compute outgoing transitions by (ilabel, olabel) pair
        // This properly handles transducers where ilabel != olabel
        let mut transitions: FxHashMap<(Label, 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 - group by (ilabel, olabel) pair to preserve both labels
            for arc in fst.arcs(state) {
                let next_weight = weight.times(&arc.weight);
                transitions
                    .entry((arc.ilabel, arc.olabel))
                    .or_insert_with(WeightedSubset::new)
                    .insert(arc.nextstate, next_weight);
            }
        }

        // set final weight if non-zero
        if !<W as num_traits::Zero>::is_zero(&final_weight) {
            result.set_final(current_state, final_weight);
        }

        // add transitions
        for ((ilabel, olabel), mut next_subset) in transitions {
            // normalize subset
            if let Some(norm_weight) = next_subset.normalize() {
                let next_state = match subset_map.get(&next_subset) {
                    Some(&state) => state,
                    None => {
                        let state = result.add_state();
                        subset_map.insert(next_subset.clone(), state);
                        queue.push((next_subset, state));
                        state
                    }
                };

                result.add_arc(
                    current_state,
                    Arc::new(ilabel, olabel, norm_weight, next_state),
                );
            }
        }
    }

    Ok(result)
}

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

    #[test]
    fn test_determinize_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(&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_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(&fst).unwrap();

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

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

    #[test]
    fn test_determinize_normalization_all_infinity() {
        // Test normalization when all weights are infinity (zero element)
        let mut subset = WeightedSubset::<TropicalWeight>::new();
        let zero = TropicalWeight::zero(); // infinity
        subset.insert(0, zero);
        subset.insert(1, zero);
        subset.insert(2, zero);

        // Normalization should return None when all weights are infinity
        let result = subset.normalize();
        assert!(
            result.is_none(),
            "Normalization should return None for all-infinity weights"
        );
    }

    #[test]
    fn test_determinize_normalization_mixed_weights() {
        // Test normalization with mixed infinity and finite weights
        let mut subset = WeightedSubset::<TropicalWeight>::new();
        subset.insert(0, TropicalWeight::new(5.0));
        subset.insert(1, TropicalWeight::new(3.0));
        subset.insert(2, TropicalWeight::zero()); // infinity

        // Normalization should work (min is 3.0, not infinity)
        let result = subset.normalize();
        assert!(
            result.is_some(),
            "Normalization should work with finite minimum"
        );

        if let Some(norm_weight) = result {
            // Normalized weights: 5.0-3.0=2.0, 3.0-3.0=0.0, infinity stays infinity
            // But actually, infinity divided by 3.0 is still infinity
            // So subset should have: 2.0, 0.0, infinity
            assert_eq!(norm_weight, TropicalWeight::new(3.0));
        }
    }

    #[test]
    fn test_determinize_normalization_empty_subset() {
        // Test normalization on empty subset
        let mut subset = WeightedSubset::<TropicalWeight>::new();
        let result = subset.normalize();
        assert!(
            result.is_none(),
            "Normalization should return None for empty subset"
        );
    }

    #[test]
    fn test_determinize_with_infinity_weights() {
        // Test determinization when FST has infinity weights
        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());

        // Add arcs with some infinity weights
        fst.add_arc(s0, Arc::new(1, 1, TropicalWeight::new(1.0), s1));
        fst.add_arc(s0, Arc::new(1, 1, TropicalWeight::zero(), s2)); // infinity weight

        // Determinization should handle this correctly
        let det: VectorFst<TropicalWeight> = determinize(&fst).unwrap();
        assert!(det.start().is_some());
        // Should still produce deterministic result
        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"
                );
            }
        }
    }

    #[test]
    fn test_determinize_transducer() {
        // Test determinization with a transducer (ilabel != olabel)
        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());

        // Same input label 'a' but different output labels 'x' and 'y'
        fst.add_arc(
            s0,
            Arc::new(b'a' as u32, b'x' as u32, TropicalWeight::new(1.0), s1),
        );
        fst.add_arc(
            s0,
            Arc::new(b'a' as u32, b'y' as u32, TropicalWeight::new(2.0), s2),
        );
        fst.add_arc(
            s1,
            Arc::new(b'b' as u32, b'z' as u32, TropicalWeight::one(), s3),
        );
        fst.add_arc(
            s2,
            Arc::new(b'b' as u32, b'z' as u32, TropicalWeight::one(), s3),
        );

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

        // Verify both output labels are preserved
        if let Some(start) = det.start() {
            let arcs: Vec<_> = det.arcs(start).collect();
            // Should have two arcs from start (one for each (ilabel, olabel) pair)
            assert_eq!(arcs.len(), 2, "Should preserve both output labels");

            // Check that we have arcs with different output labels
            let olabels: std::collections::HashSet<_> = arcs.iter().map(|a| a.olabel).collect();
            assert!(
                olabels.contains(&(b'x' as u32)),
                "Should preserve output label 'x'"
            );
            assert!(
                olabels.contains(&(b'y' as u32)),
                "Should preserve output label 'y'"
            );

            // All arcs should have same input label
            for arc in &arcs {
                assert_eq!(arc.ilabel, b'a' as u32);
            }
        }
    }
}