atlas-archive-core 1.1.0

//! High-order context mixing for Atlas (ZPAQ/PAQ8 inspired).
//! Implements multi-order probability blending for extreme ratios.

use crate::alloc::collections::BTreeMap;
use crate::alloc::vec::Vec;
/// A single context node in the high-order models.
/// A single context node in the high-order models.
#[derive(Clone, Debug)]
pub struct ContextNode {
    pub count: [u32; 256],
    pub total: u32,
}

impl Default for ContextNode {
    fn default() -> Self {
        Self {
            count: [0u32; 256],
            total: 0,
        }
    }
}

/// An n-gram model for a specific order.
#[derive(Clone, Default, Debug)]
struct NgramState {
    nodes: BTreeMap<Vec<u8>, ContextNode>,
}

/// High-order context mixer combining multiple order predictions.
pub struct ContextMixer {
    orders: Vec<usize>,
    models: Vec<NgramState>,
    /// Total nodes across all models to stay under RAM budget.
    max_nodes: usize,
    /// Current count of total nodes across all orders.
    current_nodes: usize,
    /// Prune divisor (50 = 2%, 20 = 5%, 10 = 10%)
    prune_divisor: usize,
}

impl ContextMixer {
    pub fn new(orders: Vec<usize>) -> Self {
        let mut models = Vec::with_capacity(orders.len());
        for _ in 0..orders.len() {
            models.push(NgramState::default());
        }
        Self {
            orders,
            models,
            max_nodes: 8192, // Scaled default (1024 * 8 orders approx)
            current_nodes: 0,
            prune_divisor: 20, // 5% default
        }
    }

    /// Set a custom node limit for the mixer.
    pub fn with_node_limit(mut self, limit: usize) -> Self {
        self.max_nodes = limit;
        self
    }

    /// Set a custom prune divisor.
    pub fn with_prune_divisor(mut self, divisor: usize) -> Self {
        self.prune_divisor = divisor;
        self
    }

    /// Update the mixer with a new symbol and its history.
    pub fn update(&mut self, history: &[u8], next: u8) {
        let h_len = history.len();
        for (i, &order) in self.orders.iter().enumerate() {
            if h_len >= order {
                let context = &history[h_len - order..];
                let node = self.models[i]
                    .nodes
                    .entry(context.to_vec())
                    .or_insert_with(|| {
                        self.current_nodes += 1;
                        ContextNode::default()
                    });

                node.count[next as usize] = node.count[next as usize].saturating_add(1);
                node.total = node.total.saturating_add(1);
            }
        }

        if self.current_nodes > self.max_nodes {
            self.prune();
        }
    }

    /// Prune the models to stay within node budget.
    /// Uses a "Curator" strategy: remove lowest utility nodes.
    fn prune(&mut self) {
        // We prune a fraction of the budget (default 5%) to avoid thrashing.
        let target_remove = self.max_nodes / self.prune_divisor;
        let mut removed = 0;

        // Simple strategy: Iterate and remove nodes with total < 2 in a single pass.
        // If still over, we'd need a more expensive sorted prune.
        for model in &mut self.models {
            let mut keys_to_remove = Vec::new();
            for (key, node) in &model.nodes {
                if node.total < 2 {
                    keys_to_remove.push(key.clone());
                    if keys_to_remove.len() + removed >= target_remove {
                        break;
                    }
                }
            }
            for key in keys_to_remove {
                model.nodes.remove(&key);
                removed += 1;
                self.current_nodes -= 1;
            }
            if removed >= target_remove {
                break;
            }
        }

        // Emergency fallback: If still way over, clear smallest orders
        if self.current_nodes > self.max_nodes * 11 / 10 {
            for model in &mut self.models {
                let diff = model.nodes.len();
                model.nodes.clear();
                self.current_nodes -= diff;
                if self.current_nodes <= self.max_nodes {
                    break;
                }
            }
        }
    }

    /// Query the mixer for a probability distribution.
    /// PAQ Standard: Mix predictions from multiple orders.
    #[inline]
    pub fn predict(&self, history: &[u8]) -> Option<[u64; 256]> {
        let mut mixed_probs = [0u64; 256];
        let mut total_weight = 0u64;
        let h_len = history.len();

        for (i, &order) in self.orders.iter().enumerate() {
            if h_len >= order {
                let context = &history[h_len - order..];
                if let Some(node) = self.models[i].nodes.get(context).filter(|n| n.total > 0) {
                    // Order-based weighting: higher orders get quadratic weight.
                    let weight = (order + 1) as u64 * (order + 1) as u64;
                    let node_total = node.total as u64;

                    // SIMD-friendly: process for auto-vectorization
                    let multiplier = 65536 * weight;
                    for (p, count) in mixed_probs.iter_mut().zip(node.count.iter()) {
                        *p += (*count as u64 * multiplier) / node_total;
                    }
                    total_weight += weight;
                }
            }
        }

        if total_weight == 0 {
            None
        } else {
            // SIMD-friendly normalization
            for prob in &mut mixed_probs {
                *prob /= total_weight;
            }
            Some(mixed_probs)
        }
    }
}