arcweight 0.3.0

//! OpenFST binary format compatibility layer.
//!
//! This module provides read/write support for the OpenFST binary format,
//! enabling interoperability with the widely-used OpenFST C++ library and
//! its extensive ecosystem of tools.
//!
//! # Overview
//!
//! OpenFST is the de facto standard library for weighted finite-state transducer
//! manipulation, used extensively in:
//!
//! - **Speech recognition:** Kaldi, ESPnet, wav2vec2, Whisper pipelines
//! - **Natural language processing:** Pronunciation lexicons, G2P models
//! - **Machine translation:** Phrase-based and neural hybrid systems
//! - **Text processing:** Tokenization, normalization, transliteration
//!
//! This module ensures seamless integration between ArcWeight and existing
//! OpenFST-based pipelines.
//!
//! # Format Specification
//!
//! The OpenFST binary format uses little-endian byte ordering:
//!
//! ```text
//! +------------------+------------------+
//! | Magic (4B)       | FST Type (8B)    |  Header
//! +------------------+------------------+
//! | Arc Type (16B)   | Version (4B)     |
//! +------------------+------------------+
//! | Flags (4B)       | Properties (8B)  |
//! +------------------+------------------+
//! | Start State (8B) | Num States (8B)  |
//! +------------------+------------------+
//! | Num Arcs (8B)    |                  |
//! +------------------+------------------+
//! | State 0: final_weight (4B)         |  States
//! |          num_arcs (8B)             |
//! |          arc_0 ... arc_n           |
//! +------------------------------------+
//! | State 1: ...                       |
//! +------------------------------------+
//! ```
//!
//! Each arc contains:
//! - Input label (4 bytes, i32)
//! - Output label (4 bytes, i32)
//! - Weight (4 bytes, f32 for tropical)
//! - Next state (4 bytes, i32)
//!
//! ## Header Constants
//!
//! - **Magic number:** 2125659606
//! - **FST type:** "vector" (null-padded to 8 bytes)
//! - **Arc type:** "tropical" (null-padded to 16 bytes)
//! - **Version:** 1
//!
//! # Current Limitations
//!
//! | Feature | Status |
//! |---------|--------|
//! | Vector FSTs | Supported |
//! | Const FSTs | Not supported |
//! | Tropical semiring | Supported |
//! | Log semiring | Not supported |
//! | Symbol tables | Not supported |
//! | Property flags | Written as 0 |
//!
//! # Examples
//!
//! ## Reading OpenFST Files
//!
//! ```no_run
//! use arcweight::prelude::*;
//! use arcweight::io::read_openfst;
//! use std::fs::File;
//!
//! // Read FST created by OpenFST tools (fstcompile, fstdeterminize, etc.)
//! let mut file = File::open("model.fst")?;
//! let fst: VectorFst<TropicalWeight> = read_openfst(&mut file)?;
//!
//! println!("Loaded FST with {} states", fst.num_states());
//! # Ok::<(), Box<dyn std::error::Error>>(())
//! ```
//!
//! ## Writing OpenFST Files
//!
//! ```no_run
//! use arcweight::prelude::*;
//! use arcweight::io::write_openfst;
//! use std::fs::File;
//!
//! 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(0.5), s1));
//!
//! // Write in OpenFST format
//! let mut file = File::create("output.fst")?;
//! write_openfst(&fst, &mut file)?;
//!
//! // Now usable with OpenFST command-line tools:
//! // $ fstinfo output.fst
//! // $ fstdraw output.fst | dot -Tpng > fst.png
//! // $ fstprint output.fst
//! # Ok::<(), Box<dyn std::error::Error>>(())
//! ```
//!
//! ## Pipeline Integration with Kaldi
//!
//! ```no_run
//! use arcweight::prelude::*;
//! use arcweight::io::{read_openfst, write_openfst};
//! use std::fs::File;
//!
//! // Read lexicon FST from Kaldi
//! let mut file = File::open("L.fst")?;
//! let lexicon: VectorFst<TropicalWeight> = read_openfst(&mut file)?;
//!
//! // Optimize with ArcWeight
//! let optimized: VectorFst<TropicalWeight> = minimize(&lexicon)?;
//!
//! // Write back for use in Kaldi
//! let mut out = File::create("L_opt.fst")?;
//! write_openfst(&optimized, &mut out)?;
//! # Ok::<(), Box<dyn std::error::Error>>(())
//! ```
//!
//! # Weight Representation
//!
//! The tropical semiring in OpenFST uses IEEE 754 single-precision floats:
//!
//! | Semiring Element | OpenFST Representation |
//! |------------------|------------------------|
//! | Zero (annihilator) | `+inf` (positive infinity) |
//! | One (identity) | `0.0` |
//! | General weight | float value (lower = better) |
//!
//! # State Numbering
//!
//! - States are numbered consecutively from 0
//! - State IDs are preserved exactly during read/write
//! - The special value -1 indicates "no start state"
//!
//! # References
//!
//! - Allauzen, C., Riley, M., Schalkwyk, J., Skut, W., & Mohri, M. (2007).
//!   OpenFst: A general and efficient weighted finite-state transducer library.
//!   In *Implementation and Application of Automata* (CIAA 2007), pp. 11-23.
//!   Springer. <https://doi.org/10.1007/978-3-540-76336-9_3>
//!
//! - OpenFST Library. <https://www.openfst.org/>
//!
//! - Povey, D., et al. (2011). The Kaldi speech recognition toolkit.
//!   In *IEEE Workshop on Automatic Speech Recognition and Understanding*.

use crate::arc::Arc;
use crate::fst::{Fst, MutableFst, StateId};
use crate::semiring::{Semiring, TropicalWeight};
use crate::{Error, Result};
use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
use std::io::{Read, Seek, Write};

const OPENFST_MAGIC: i32 = 2_125_659_606;

/// FST type in OpenFST format
#[derive(Debug, Clone, Copy)]
#[repr(i32)]
#[allow(dead_code)]
enum FstType {
    Vector = 1,
    Const = 2,
}

/// Writes an FST in OpenFST binary format.
///
/// Serializes the FST to the standard OpenFST binary format, making it compatible
/// with OpenFST command-line tools (`fstinfo`, `fstdraw`, `fstprint`, etc.) and
/// other libraries that read OpenFST files (Kaldi, ESPnet, etc.).
///
/// # Type Parameters
///
/// * `F` - The FST type, must implement `Fst<TropicalWeight>`
/// * `Writer` - The output writer type, must implement `Write + Seek`
///
/// # Arguments
///
/// * `fst` - Reference to the FST to serialize
/// * `writer` - Mutable reference to the output writer (must support seeking)
///
/// # Returns
///
/// Returns `Ok(())` on successful serialization.
///
/// # Errors
///
/// Returns [`Error::Io`] if the writer encounters an I/O error
/// during writing or seeking.
///
/// # Complexity
///
/// - **Time:** O(|V| + |E|) where |V| is the number of states and |E| is
///   the number of arcs
/// - **Space:** O(1) additional space (streaming write)
///
/// # Format Details
///
/// The output file uses:
/// - Magic number: 2125659606
/// - FST type: "vector"
/// - Arc type: "tropical"
/// - Little-endian byte order
/// - IEEE 754 f32 for weights (+inf = zero element)
///
/// # Examples
///
/// ```no_run
/// use arcweight::prelude::*;
/// use arcweight::io::write_openfst;
/// use std::fs::File;
///
/// 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, 2, TropicalWeight::new(0.5), s1));
///
/// let mut file = File::create("output.fst")?;
/// write_openfst(&fst, &mut file)?;
/// # Ok::<(), Box<dyn std::error::Error>>(())
/// ```
///
/// # See Also
///
/// - [`read_openfst`] - Read FSTs in OpenFST format
/// - [`write_binary`](super::write_binary) - Native binary format (Rust-only)
pub fn write_openfst<F, Writer>(fst: &F, writer: &mut Writer) -> Result<()>
where
    F: Fst<TropicalWeight>,
    Writer: Write + Seek,
{
    // write header
    writer.write_i32::<LittleEndian>(OPENFST_MAGIC)?;
    writer.write_all(b"vector\0\0")?; // fst type (8 bytes)
    writer.write_all(b"tropical\0\0\0\0\0\0\0\0")?; // arc type (16 bytes)
    writer.write_i32::<LittleEndian>(1)?; // version
    writer.write_i32::<LittleEndian>(0)?; // flags
    writer.write_i64::<LittleEndian>(0)?; // properties
    writer.write_i64::<LittleEndian>(fst.start().unwrap_or(u32::MAX) as i64)?;
    writer.write_i64::<LittleEndian>(fst.num_states() as i64)?;

    // Count total arcs first
    let total_arcs = fst.states().map(|s| fst.num_arcs(s) as i64).sum::<i64>();
    writer.write_i64::<LittleEndian>(total_arcs)?;

    // write states
    for state in fst.states() {
        let final_weight = fst
            .final_weight(state)
            .map(|w| *w.value())
            .unwrap_or(f32::INFINITY);

        writer.write_f32::<LittleEndian>(final_weight)?;
        writer.write_i64::<LittleEndian>(fst.num_arcs(state) as i64)?;

        for arc in fst.arcs(state) {
            writer.write_i32::<LittleEndian>(arc.ilabel as i32)?;
            writer.write_i32::<LittleEndian>(arc.olabel as i32)?;
            writer.write_f32::<LittleEndian>(*arc.weight.value())?;
            writer.write_i32::<LittleEndian>(arc.nextstate as i32)?;
        }
    }

    Ok(())
}

/// Reads an FST from OpenFST binary format.
///
/// Parses the standard OpenFST binary format, enabling loading of FSTs created
/// by OpenFST command-line tools (`fstcompile`, `fstdeterminize`, etc.) or other
/// compatible libraries.
///
/// # Type Parameters
///
/// * `M` - The mutable FST type to construct, must implement
///   `MutableFst<TropicalWeight> + Default`
/// * `Reader` - The input reader type, must implement `Read`
///
/// # Arguments
///
/// * `reader` - Mutable reference to the input reader
///
/// # Returns
///
/// Returns the parsed FST on success.
///
/// # Errors
///
/// Returns [`Error::Serialization`] if the magic
/// number is invalid (file is not in OpenFST format).
///
/// Returns [`Error::Io`] if the reader encounters an I/O error.
///
/// # Complexity
///
/// - **Time:** O(|V| + |E|) where |V| is the number of states and |E| is
///   the number of arcs
/// - **Space:** O(|V| + |E|) for the output FST
///
/// # Correctness
///
/// This function guarantees compatibility with OpenFST:
/// - Reads files written by OpenFST tools
/// - Produces FSTs that can be written back in OpenFST format
/// - Language is preserved: L(read(write(T))) = L(T)
/// - State IDs are preserved exactly
///
/// # Limitations
///
/// - Only supports vector FSTs (not const FSTs)
/// - Only supports tropical semiring weights
/// - Does not read symbol tables (use separate `.syms` files)
/// - Does not preserve property flags
///
/// # Examples
///
/// ```no_run
/// use arcweight::prelude::*;
/// use arcweight::io::read_openfst;
/// use std::fs::File;
///
/// let mut file = File::open("model.fst")?;
/// let fst: VectorFst<TropicalWeight> = read_openfst(&mut file)?;
///
/// println!("States: {}", fst.num_states());
/// println!("Start: {:?}", fst.start());
/// # Ok::<(), Box<dyn std::error::Error>>(())
/// ```
///
/// # See Also
///
/// - [`write_openfst`] - Write FSTs in OpenFST format
/// - [`read_binary`](super::read_binary) - Native binary format (faster, Rust-only)
pub fn read_openfst<M, Reader>(reader: &mut Reader) -> Result<M>
where
    M: MutableFst<TropicalWeight> + Default,
    Reader: Read,
{
    // read header
    let magic = reader.read_i32::<LittleEndian>()?;
    if magic != OPENFST_MAGIC {
        return Err(Error::Serialization("Invalid OpenFST magic".into()));
    }

    // skip fst type and arc type
    let mut buf = [0u8; 24];
    reader.read_exact(&mut buf)?;

    let _version = reader.read_i32::<LittleEndian>()?;
    let _flags = reader.read_i32::<LittleEndian>()?;
    let _properties = reader.read_i64::<LittleEndian>()?;
    let start = reader.read_i64::<LittleEndian>()?;
    let num_states = reader.read_i64::<LittleEndian>()? as usize;
    let _num_arcs = reader.read_i64::<LittleEndian>()?;

    let mut fst = M::default();

    // create states
    for _ in 0..num_states {
        fst.add_state();
    }

    // set start
    if start >= 0 {
        fst.set_start(start as StateId);
    }

    // read states
    for state in 0..num_states {
        let final_weight = reader.read_f32::<LittleEndian>()?;
        if final_weight != f32::INFINITY {
            fst.set_final(state as StateId, TropicalWeight::new(final_weight));
        }

        let num_arcs = reader.read_i64::<LittleEndian>()? as usize;
        for _ in 0..num_arcs {
            let ilabel = reader.read_i32::<LittleEndian>()? as u32;
            let olabel = reader.read_i32::<LittleEndian>()? as u32;
            let weight = reader.read_f32::<LittleEndian>()?;
            let nextstate = reader.read_i32::<LittleEndian>()? as u32;

            fst.add_arc(
                state as StateId,
                Arc::new(ilabel, olabel, TropicalWeight::new(weight), nextstate),
            );
        }
    }

    Ok(fst)
}

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

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

        fst.set_start(s0);
        fst.set_final(s1, TropicalWeight::new(1.5));
        fst.add_arc(s0, Arc::new(1, 2, TropicalWeight::new(0.5), s1));

        let buffer = Vec::new();
        let mut cursor = Cursor::new(buffer);
        write_openfst(&fst, &mut cursor).unwrap();

        cursor.set_position(0);
        let read_fst: VectorFst<TropicalWeight> =
            read_openfst::<VectorFst<TropicalWeight>, _>(&mut cursor).unwrap();

        assert_eq!(read_fst.num_states(), fst.num_states());
        assert_eq!(read_fst.start(), fst.start());
        assert_eq!(read_fst.num_arcs_total(), fst.num_arcs_total());
    }

    #[test]
    fn test_openfst_empty_fst() {
        let fst = VectorFst::<TropicalWeight>::new();

        let buffer = Vec::new();
        let mut cursor = Cursor::new(buffer);
        write_openfst(&fst, &mut cursor).unwrap();

        cursor.set_position(0);
        let read_fst: VectorFst<TropicalWeight> =
            read_openfst::<VectorFst<TropicalWeight>, _>(&mut cursor).unwrap();

        assert!(read_fst.is_empty());
    }

    #[test]
    fn test_openfst_compatibility() {
        // Test that our format can handle typical OpenFST constructs
        let mut fst = VectorFst::<TropicalWeight>::new();

        // Create a more complex FST structure
        for _i in 0..10 {
            fst.add_state();
        }

        fst.set_start(0);
        fst.set_final(9, TropicalWeight::new(0.0));

        // Add various arc types
        fst.add_arc(0, Arc::new(1, 1, TropicalWeight::new(1.0), 1));
        fst.add_arc(1, Arc::epsilon(TropicalWeight::new(0.0), 2));
        fst.add_arc(2, Arc::new(2, 3, TropicalWeight::new(2.0), 3));

        // Add some parallel arcs
        fst.add_arc(0, Arc::new(4, 4, TropicalWeight::new(3.0), 4));
        fst.add_arc(4, Arc::new(5, 5, TropicalWeight::new(1.0), 9));

        let buffer = Vec::new();
        let mut cursor = Cursor::new(buffer);
        write_openfst(&fst, &mut cursor).unwrap();

        cursor.set_position(0);
        let read_fst: VectorFst<TropicalWeight> =
            read_openfst::<VectorFst<TropicalWeight>, _>(&mut cursor).unwrap();

        assert_eq!(read_fst.num_states(), fst.num_states());
        assert_eq!(read_fst.start(), fst.start());

        // Verify specific structural elements
        assert!(read_fst.is_final(9));
        assert_eq!(read_fst.num_arcs(0), 2); // Two outgoing arcs from start
    }
}