zipora 2.1.5

//! Directed Acyclic Word Graph (DAWG) and Deterministic Finite Automaton (DFA) implementations
//!
//! This module provides high-performance DAWG and DFA structures with advanced optimization
//! techniques including nested trie DAWG, rank-select integration, and compressed storage.

use crate::error::{Result, ZiporaError};
use crate::fsa::cache::{FsaCache, FsaCacheConfig};
use crate::fsa::traits::{FiniteStateAutomaton, Trie, TrieStats, StatisticsProvider};
use crate::StateId;
use crate::memory::SecureMemoryPool;
use crate::succinct::rank_select::{RankSelectInterleaved256, RankSelectOps};
use crate::succinct::BitVector;
use std::collections::HashMap;
use std::sync::Arc;

/// Configuration for DAWG construction
#[derive(Debug, Clone)]
pub struct DawgConfig {
    /// Use rank-select acceleration for terminal states
    pub use_rank_select: bool,
    /// Enable FSA caching for improved lookup performance  
    pub enable_cache: bool,
    /// Cache configuration
    pub cache_config: FsaCacheConfig,
    /// Maximum number of states in DAWG
    pub max_states: usize,
    /// Enable compressed storage
    pub compressed_storage: bool,
}

impl Default for DawgConfig {
    fn default() -> Self {
        Self {
            use_rank_select: true,
            enable_cache: true,
            cache_config: FsaCacheConfig::default(),
            max_states: 1_000_000,
            compressed_storage: true,
        }
    }
}

impl DawgConfig {
    /// Create configuration optimized for memory efficiency
    pub fn memory_efficient() -> Self {
        Self {
            cache_config: FsaCacheConfig::memory_efficient(),
            max_states: 100_000,
            compressed_storage: true,
            ..Default::default()
        }
    }

    /// Create configuration optimized for query performance
    pub fn performance_optimized() -> Self {
        Self {
            cache_config: FsaCacheConfig::large(),
            max_states: 10_000_000,
            use_rank_select: true,
            ..Default::default()
        }
    }
}

/// Terminal state handling strategy
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum TerminalStrategy {
    /// Store terminal flags in a bit vector with rank-select support
    RankSelect,
    /// Store terminal flags in a simple bit vector
    BitVector,
    /// Store terminal flags inline with states
    Inline,
}

/// State in the DAWG/DFA (8-byte optimized representation)
#[derive(Debug, Clone, Copy)]
#[repr(C)]
pub struct DawgState {
    /// Base index for child transitions
    pub child_base: u32,
    /// State flags and metadata (24-bit parent + 8-bit flags)
    pub flags_and_parent: u32,
}

impl DawgState {
    /// Create a new DAWG state
    pub fn new(child_base: u32, parent: u32, is_terminal: bool, is_final: bool) -> Self {
        let flags = if is_terminal { 0x80 } else { 0 } | if is_final { 0x40 } else { 0 };
        let flags_and_parent = (parent & 0x00FFFFFF) | ((flags as u32) << 24);
        
        Self {
            child_base,
            flags_and_parent,
        }
    }

    /// Get parent state ID
    pub fn parent(&self) -> u32 {
        self.flags_and_parent & 0x00FFFFFF
    }

    /// Check if this is a terminal state
    #[inline]
    pub fn is_terminal(&self) -> bool {
        (self.flags_and_parent & 0x80000000) != 0
    }

    /// Check if this is a final state
    #[inline]
    pub fn is_final(&self) -> bool {
        (self.flags_and_parent & 0x40000000) != 0
    }

    /// Get flags byte
    pub fn flags(&self) -> u8 {
        ((self.flags_and_parent >> 24) & 0xFF) as u8
    }

    /// Set terminal flag
    pub fn set_terminal(&mut self, is_terminal: bool) {
        if is_terminal {
            self.flags_and_parent |= 0x80000000;
        } else {
            self.flags_and_parent &= !0x80000000;
        }
    }

    /// Set final flag
    pub fn set_final(&mut self, is_final: bool) {
        if is_final {
            self.flags_and_parent |= 0x40000000;
        } else {
            self.flags_and_parent &= !0x40000000;
        }
    }
}

/// Transition table for DAWG states.
///
/// Uses sparse HashMap storage only — the dense state*256 approach was removed
/// because it wastes O(256 * states) = massive memory for typical FSAs where
/// average fan-out is 4-8. HashMap gives O(transitions) memory.
#[derive(Debug, Clone)]
pub struct TransitionTable {
    /// Sparse transition map: (from_state, symbol) -> to_state
    pub sparse_table: HashMap<(u32, u8), u32>,
    /// Number of states
    pub num_states: u32,
}

impl TransitionTable {
    /// Create a new transition table.
    pub fn new(num_states: u32, _use_dense: bool) -> Self {
        Self {
            sparse_table: HashMap::new(),
            num_states,
        }
    }

    /// Add a transition.
    #[inline]
    pub fn add_transition(&mut self, from_state: u32, symbol: u8, to_state: u32) -> Result<()> {
        self.sparse_table.insert((from_state, symbol), to_state);
        Ok(())
    }

    /// Get transition target state.
    #[inline]
    pub fn get_transition(&self, from_state: u32, symbol: u8) -> Option<u32> {
        self.sparse_table.get(&(from_state, symbol)).copied()
    }

    /// Get all outgoing transitions from a state (sorted by symbol).
    pub fn get_outgoing_transitions(&self, from_state: u32) -> Vec<(u8, u32)> {
        let mut transitions: Vec<(u8, u32)> = self.sparse_table.iter()
            .filter_map(|(&(state, symbol), &target)| {
                if state == from_state { Some((symbol, target)) } else { None }
            })
            .collect();
        transitions.sort_unstable_by_key(|(symbol, _)| *symbol);
        transitions
    }

    /// Estimate memory usage in bytes.
    #[inline]
    pub fn memory_usage(&self) -> usize {
        self.sparse_table.len() * (std::mem::size_of::<(u32, u8)>() + std::mem::size_of::<u32>())
    }
}

/// Nested Trie DAWG with rank-select optimization
pub struct NestedTrieDawg {
    /// DAWG configuration
    config: DawgConfig,
    /// DAWG states
    states: Vec<DawgState>,
    /// Transition table
    transitions: TransitionTable,
    /// Terminal state handling
    terminal_strategy: TerminalStrategy,
    /// Terminal states bit vector
    terminal_bits: Option<BitVector>,
    /// Rank-select structure for terminal states
    terminal_rank_select: Option<RankSelectInterleaved256>,
    /// FSA cache for performance optimization
    cache: Option<FsaCache>,
    /// Root state ID
    root_state: u32,
    /// Number of keys stored
    num_keys: usize,
    /// Memory pool for allocation
    memory_pool: Option<Arc<SecureMemoryPool>>,
}

impl NestedTrieDawg {
    /// Create a new nested trie DAWG
    pub fn new() -> Result<Self> {
        Self::with_config(DawgConfig::default())
    }

    /// Create a new nested trie DAWG with configuration
    pub fn with_config(config: DawgConfig) -> Result<Self> {
        let cache = if config.enable_cache {
            Some(FsaCache::with_config(config.cache_config.clone())?)
        } else {
            None
        };

        let memory_pool = Some(SecureMemoryPool::new(
            crate::memory::SecurePoolConfig::small_secure()
        )?);

        let transitions = TransitionTable::new(config.max_states as u32, false);

        let terminal_strategy = if config.use_rank_select {
            TerminalStrategy::RankSelect
        } else {
            TerminalStrategy::BitVector
        };

        Ok(Self {
            config,
            states: Vec::new(),
            transitions,
            terminal_strategy,
            terminal_bits: None,
            terminal_rank_select: None,
            cache,
            root_state: 0,
            num_keys: 0,
            memory_pool,
        })
    }

    /// Build DAWG from a set of keys
    pub fn build_from_keys<I, K>(&mut self, keys: I) -> Result<()>
    where
        I: IntoIterator<Item = K>,
        K: AsRef<[u8]>,
    {
        // Clear existing data
        self.clear();

        // Create root state
        self.root_state = self.add_state(0, false, false)?;

        // Build trie structure first
        for key in keys {
            self.insert_key(key.as_ref())?;
        }

        // Convert to DAWG by merging equivalent states
        self.convert_to_dawg()?;

        // Build terminal rank-select structure if needed
        if self.terminal_strategy == TerminalStrategy::RankSelect {
            self.build_terminal_rank_select()?;
        }

        Ok(())
    }

    /// Insert a single key into the trie structure
    fn insert_key(&mut self, key: &[u8]) -> Result<()> {
        let mut current_state = self.root_state;

        // Traverse/create path for the key
        for &symbol in key {
            if let Some(next_state) = self.transitions.get_transition(current_state, symbol) {
                current_state = next_state;
            } else {
                // Create new state
                let new_state = self.add_state(current_state, false, false)?;
                self.transitions.add_transition(current_state, symbol, new_state)?;
                current_state = new_state;
            }
        }

        // Mark final state as terminal
        if (current_state as usize) < self.states.len() {
            self.states[current_state as usize].set_terminal(true);
            self.num_keys += 1;
        }

        Ok(())
    }

    /// Add a new state to the DAWG
    fn add_state(&mut self, parent: u32, is_terminal: bool, is_final: bool) -> Result<u32> {
        let state_id = self.states.len() as u32;
        
        if state_id >= self.config.max_states as u32 {
            return Err(ZiporaError::invalid_data("Maximum states exceeded"));
        }

        let state = DawgState::new(0, parent, is_terminal, is_final);
        self.states.push(state);

        // Cache the state if caching is enabled
        if let Some(ref mut cache) = self.cache {
            cache.cache_state(parent, 0, is_terminal)?;
        }

        Ok(state_id)
    }

    /// Convert trie to DAWG by merging equivalent states
    fn convert_to_dawg(&mut self) -> Result<()> {
        // Signature-based state equivalence
        let mut state_signatures: HashMap<Vec<u8>, u32> = HashMap::new();
        let mut state_mapping: HashMap<u32, u32> = HashMap::new();

        // Process states bottom-up (reverse topological order)
        for state_id in (0..self.states.len()).rev() {
            let signature = self.compute_state_signature(state_id as u32)?;
            
            if let Some(&equivalent_state) = state_signatures.get(&signature) {
                // Found equivalent state - merge
                state_mapping.insert(state_id as u32, equivalent_state);
            } else {
                // New unique state
                state_signatures.insert(signature, state_id as u32);
                state_mapping.insert(state_id as u32, state_id as u32);
            }
        }

        // Update transitions based on state mapping
        self.remap_transitions(&state_mapping)?;

        // Compact state array by removing merged states
        self.compact_states(&state_mapping)?;

        Ok(())
    }

    /// Compute signature for state equivalence
    fn compute_state_signature(&self, state_id: u32) -> Result<Vec<u8>> {
        let mut signature = Vec::new();

        // Add terminal flag
        signature.push(if self.states[state_id as usize].is_terminal() { 1 } else { 0 });

        // Add sorted outgoing transitions
        let mut transitions = self.transitions.get_outgoing_transitions(state_id);
        transitions.sort_by_key(|(symbol, _)| *symbol);

        for (symbol, target_state) in transitions {
            signature.push(symbol);
            signature.extend_from_slice(&target_state.to_le_bytes());
        }

        Ok(signature)
    }

    /// Remap transitions based on state mapping.
    fn remap_transitions(&mut self, state_mapping: &HashMap<u32, u32>) -> Result<()> {
        let old_transitions = std::mem::replace(
            &mut self.transitions,
            TransitionTable::new(self.states.len() as u32, false)
        );

        for (&(from_state, symbol), &to_state) in &old_transitions.sparse_table {
            if let (Some(&new_from), Some(&new_to)) =
                (state_mapping.get(&from_state), state_mapping.get(&to_state)) {
                self.transitions.add_transition(new_from, symbol, new_to)?;
            }
        }

        Ok(())
    }

    /// Compact states array after merging.
    fn compact_states(&mut self, state_mapping: &HashMap<u32, u32>) -> Result<()> {
        let mut new_states = Vec::new();
        let mut compaction_mapping: HashMap<u32, u32> = HashMap::new();

        let mut new_id = 0u32;
        for old_id in 0..self.states.len() as u32 {
            if let Some(&mapped_id) = state_mapping.get(&old_id) {
                if mapped_id == old_id {
                    compaction_mapping.insert(old_id, new_id);
                    new_states.push(self.states[old_id as usize]);
                    new_id += 1;
                }
            }
        }

        let old_transitions = std::mem::replace(
            &mut self.transitions,
            TransitionTable::new(new_states.len() as u32, false)
        );

        for (&(from_state, symbol), &to_state) in &old_transitions.sparse_table {
            if let (Some(&compact_from), Some(&compact_to)) =
                (compaction_mapping.get(&from_state), compaction_mapping.get(&to_state)) {
                self.transitions.add_transition(compact_from, symbol, compact_to)?;
            }
        }

        self.root_state = compaction_mapping.get(&self.root_state).copied().unwrap_or(0);
        self.states = new_states;

        Ok(())
    }

    /// Build rank-select structure for terminal states
    fn build_terminal_rank_select(&mut self) -> Result<()> {
        let mut terminal_bits = BitVector::new();

        for state in &self.states {
            terminal_bits.push(state.is_terminal())?;
        }

        self.terminal_rank_select = Some(RankSelectInterleaved256::new(terminal_bits.clone())?);
        self.terminal_bits = Some(terminal_bits);

        Ok(())
    }

    /// Clear all data
    pub fn clear(&mut self) {
        self.states.clear();
        self.transitions = TransitionTable::new(self.config.max_states as u32, false);
        self.terminal_bits = None;
        self.terminal_rank_select = None;
        self.root_state = 0;
        self.num_keys = 0;

        if let Some(ref mut cache) = self.cache {
            cache.clear();
        }
    }

    /// Map a terminal state to its word index (dictionary rank).
    /// Uses rank-select for O(1) lookup.
    ///
    /// Returns None if the state is not terminal.
    #[inline]
    pub fn state_to_word_id(&self, state: u32) -> Option<usize> {
        let idx = state as usize;
        if idx >= self.states.len() || !self.states[idx].is_terminal() {
            return None;
        }
        if let Some(ref rs) = self.terminal_rank_select {
            Some(rs.rank1(idx))
        } else {
            // Fallback: linear count
            Some(self.states[..=idx].iter().filter(|s| s.is_terminal()).count() - 1)
        }
    }

    /// Total number of states.
    #[inline]
    pub fn total_states(&self) -> usize {
        self.states.len()
    }

    /// Get DAWG statistics
    pub fn statistics(&self) -> DawgStats {
        let transition_memory = self.transitions.memory_usage();
        let state_memory = self.states.len() * std::mem::size_of::<DawgState>();
        let terminal_memory = self.terminal_bits.as_ref()
            .map(|_bv| std::mem::size_of::<BitVector>())
            .unwrap_or(0) + 
            self.terminal_rank_select.as_ref()
                .map(|_rs| std::mem::size_of::<RankSelectInterleaved256>())
                .unwrap_or(0);

        DawgStats {
            num_states: self.states.len(),
            num_transitions: self.transitions.sparse_table.len(),
            num_keys: self.num_keys,
            memory_usage: state_memory + transition_memory + terminal_memory,
            compression_ratio: if self.num_keys > 0 {
                self.states.len() as f64 / self.num_keys as f64
            } else {
                0.0
            },
            cache_hit_ratio: self.cache.as_ref()
                .map(|c| c.stats().hit_ratio())
                .unwrap_or(0.0),
        }
    }
}

/// Statistics for DAWG performance analysis
#[derive(Debug, Clone)]
pub struct DawgStats {
    /// Number of states in DAWG
    pub num_states: usize,
    /// Number of transitions
    pub num_transitions: usize,
    /// Number of keys stored
    pub num_keys: usize,
    /// Total memory usage in bytes
    pub memory_usage: usize,
    /// Compression ratio (states/keys)
    pub compression_ratio: f64,
    /// Cache hit ratio
    pub cache_hit_ratio: f64,
}

impl FiniteStateAutomaton for NestedTrieDawg {
    fn root(&self) -> StateId {
        self.root_state as StateId
    }

    fn is_final(&self, state: StateId) -> bool {
        if (state as usize) < self.states.len() {
            self.states[state as usize].is_terminal()
        } else {
            false
        }
    }

    fn transition(&self, state: StateId, symbol: u8) -> Option<StateId> {
        self.transitions.get_transition(state as u32, symbol).map(|s| s as StateId)
    }

    fn transitions(&self, state: StateId) -> Vec<(u8, StateId)> {
        self.transitions.get_outgoing_transitions(state as u32)
            .into_iter()
            .map(|(symbol, target)| (symbol, target as StateId))
            .collect()
    }
}

impl Trie for NestedTrieDawg {
    fn insert(&mut self, key: &[u8]) -> Result<StateId> {
        let _old_len = self.num_keys;
        self.insert_key(key)?;
        
        // Return the state ID for the inserted key
        let mut current_state = self.root_state;
        for &symbol in key {
            if let Some(next_state) = self.transitions.get_transition(current_state, symbol) {
                current_state = next_state;
            } else {
                // This shouldn't happen since we just inserted
                return Err(ZiporaError::invalid_data("Failed to find inserted key"));
            }
        }
        
        Ok(current_state as StateId)
    }

    fn len(&self) -> usize {
        self.num_keys
    }

    fn contains(&self, key: &[u8]) -> bool {
        let mut current_state = self.root_state;
        
        for &symbol in key {
            if let Some(next_state) = self.transitions.get_transition(current_state, symbol) {
                current_state = next_state;
            } else {
                return false;
            }
        }
        
        // Check if the final state is terminal
        (current_state as usize) < self.states.len() && self.states[current_state as usize].is_terminal()
    }

    fn is_empty(&self) -> bool {
        self.num_keys == 0
    }
}

impl StatisticsProvider for NestedTrieDawg {
    fn stats(&self) -> TrieStats {
        let dawg_stats = self.statistics();
        
        TrieStats {
            num_states: dawg_stats.num_states,
            num_keys: dawg_stats.num_keys,
            num_transitions: dawg_stats.num_transitions,
            max_depth: 0, // Would need tree traversal to calculate accurately
            avg_depth: 0.0, // Would need tree traversal to calculate accurately
            memory_usage: dawg_stats.memory_usage,
            bits_per_key: if dawg_stats.num_keys > 0 {
                (dawg_stats.memory_usage * 8) as f64 / dawg_stats.num_keys as f64
            } else {
                0.0
            },
        }
    }
}

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

    #[test]
    fn test_dawg_state_creation() {
        let state = DawgState::new(100, 50, true, false);
        assert_eq!(state.child_base, 100);
        assert_eq!(state.parent(), 50);
        assert!(state.is_terminal());
        assert!(!state.is_final());
    }

    #[test]
    fn test_dawg_state_flags() {
        let mut state = DawgState::new(100, 50, false, false);
        assert!(!state.is_terminal());
        assert!(!state.is_final());
        
        state.set_terminal(true);
        assert!(state.is_terminal());
        
        state.set_final(true);
        assert!(state.is_final());
    }

    #[test]
    fn test_transition_table_dense() {
        let mut table = TransitionTable::new(10, true);
        
        table.add_transition(0, b'a', 1).unwrap();
        table.add_transition(0, b'b', 2).unwrap();
        
        assert_eq!(table.get_transition(0, b'a'), Some(1));
        assert_eq!(table.get_transition(0, b'b'), Some(2));
        assert_eq!(table.get_transition(0, b'c'), None);
        
        let transitions = table.get_outgoing_transitions(0);
        assert_eq!(transitions.len(), 2);
        assert!(transitions.contains(&(b'a', 1)));
        assert!(transitions.contains(&(b'b', 2)));
    }

    #[test]
    fn test_transition_table_sparse() {
        let mut table = TransitionTable::new(10, false);
        
        table.add_transition(0, b'a', 1).unwrap();
        table.add_transition(0, b'b', 2).unwrap();
        
        assert_eq!(table.get_transition(0, b'a'), Some(1));
        assert_eq!(table.get_transition(0, b'b'), Some(2));
        assert_eq!(table.get_transition(0, b'c'), None);
        
        let transitions = table.get_outgoing_transitions(0);
        assert_eq!(transitions.len(), 2);
        assert!(transitions.contains(&(b'a', 1)));
        assert!(transitions.contains(&(b'b', 2)));
    }

    #[test]
    fn test_nested_trie_dawg_basic() {
        let mut dawg = NestedTrieDawg::new().unwrap();
        
        // Build from simple keys
        let keys = vec!["cat".as_bytes(), "car".as_bytes(), "card".as_bytes(), "care".as_bytes(), "careful".as_bytes()];
        dawg.build_from_keys(keys).unwrap();
        
        // Test containment
        assert!(dawg.contains(b"cat"));
        assert!(dawg.contains(b"car"));
        assert!(dawg.contains(b"card"));
        assert!(dawg.contains(b"care"));
        assert!(dawg.contains(b"careful"));
        assert!(!dawg.contains(b"dog"));
        assert!(!dawg.contains(b"ca"));
        
        // Test statistics
        let stats = dawg.statistics();
        assert_eq!(stats.num_keys, 5);
        assert!(stats.num_states > 0);
        assert!(stats.memory_usage > 0);
    }

    #[test]
    fn test_nested_trie_dawg_prefix_search() {
        let mut dawg = NestedTrieDawg::new().unwrap();
        
        let keys = vec!["computer".as_bytes(), "computation".as_bytes(), "compute".as_bytes(), "computing".as_bytes()];
        dawg.build_from_keys(keys).unwrap();
        
        // Test prefix search - TODO: implement prefix_search method
        // let results = dawg.prefix_search(b"comput").unwrap();
        // assert_eq!(results.len(), 4);
        
        // Test basic contains functionality for now
        assert!(dawg.contains(b"computer"));
        assert!(dawg.contains(b"computation"));
        assert!(dawg.contains(b"compute"));
        assert!(dawg.contains(b"computing"));
    }

    #[test]
    fn test_nested_trie_dawg_longest_prefix() {
        let mut dawg = NestedTrieDawg::new().unwrap();
        
        let keys = vec!["app".as_bytes(), "apple".as_bytes(), "application".as_bytes()];
        dawg.build_from_keys(keys).unwrap();
        
        assert_eq!(dawg.longest_prefix(b"app"), Some(3));
        assert_eq!(dawg.longest_prefix(b"apple"), Some(5));
        assert_eq!(dawg.longest_prefix(b"application"), Some(11));
        assert_eq!(dawg.longest_prefix(b"applications"), Some(11));
        assert_eq!(dawg.longest_prefix(b"ap"), None);
    }

    #[test]
    fn test_dawg_compression() {
        let mut dawg = NestedTrieDawg::new().unwrap();
        
        // Keys with shared suffixes should compress well in DAWG
        let keys = vec![
            "ending".as_bytes(), "reading".as_bytes(), "heading".as_bytes(), "leading".as_bytes(),
            "sending".as_bytes(), "bending".as_bytes(), "pending".as_bytes(), "mending".as_bytes()
        ];
        dawg.build_from_keys(keys).unwrap();
        
        let stats = dawg.statistics();
        assert_eq!(stats.num_keys, 8);
        
        // DAWG should have fewer states than a trie due to suffix sharing
        // (exact number depends on implementation details)
        assert!(stats.num_states < stats.num_keys * 7); // Rough upper bound
        assert!(stats.compression_ratio < 8.0); // Should be compressed
    }

    #[test]
    fn test_dawg_config_variants() {
        let memory_config = DawgConfig::memory_efficient();
        let performance_config = DawgConfig::performance_optimized();
        
        assert!(memory_config.max_states < performance_config.max_states);
        assert!(memory_config.compressed_storage);
    }

    #[test]
    fn test_dawg_empty_and_clear() {
        let mut dawg = NestedTrieDawg::new().unwrap();
        assert!(dawg.is_empty());
        assert_eq!(dawg.len(), 0);
        
        dawg.build_from_keys(vec![b"test"]).unwrap();
        assert!(!dawg.is_empty());
        assert_eq!(dawg.len(), 1);
        
        dawg.clear();
        assert!(dawg.is_empty());
        assert_eq!(dawg.len(), 0);
    }
}