srcdmp 0.1.0-alpha.1

use crate::raw::codec::freq::FREQUENCIES;
use std::cmp::Ordering;
use std::collections::{BinaryHeap, HashMap};
use std::sync::LazyLock;

/// The static Huffman tree built from [`FREQUENCIES`].
pub static HUFFMAN_TREE: LazyLock<HuffmanNode> = LazyLock::new(build_tree);

/// The static mapping from characters to their Huffman codes.
///
/// Each entry maps a `char` to a tuple of `(code_bits, bit_length)`.
pub static HUFFMAN_CODES: LazyLock<HashMap<char, (u64, u8)>> =
    LazyLock::new(|| build_codes(&HUFFMAN_TREE));

/// A node in the Huffman tree, either a leaf with a character or an
/// internal node with left and right children.
///
/// The tree is built from the empirical character frequency table and
/// used by the SDMP-C4 codec for compression.
///
/// # Examples
///
/// ```rust
/// use srcdmp::raw::codec::huffman::HUFFMAN_TREE;
///
/// // The tree is lazily initialized on first access.
/// let _tree = &*HUFFMAN_TREE;
/// ```
#[derive(Eq, PartialEq)]
pub enum HuffmanNode {
    /// A leaf node representing a single character.
    Leaf {
        /// The character stored at this leaf.
        ch: char,
        /// Frequency of this character in the training corpus.
        freq: u64,
    },
    /// An internal node combining two subtrees.
    Internal {
        /// Sum of frequencies of all characters in this subtree.
        freq: u64,
        /// Left child (associated with bit `0`).
        left: Box<Self>,
        /// Right child (associated with bit `1`).
        right: Box<Self>,
    },
}

impl HuffmanNode {
    /// Return the aggregate frequency for this node.
    #[must_use]
    pub const fn freq(&self) -> u64 {
        match self {
            Self::Leaf { freq, .. } => *freq,
            Self::Internal { freq, .. } => *freq,
        }
    }
}

impl Ord for HuffmanNode {
    fn cmp(&self, other: &Self) -> Ordering {
        other.freq().cmp(&self.freq())
    }
}

impl PartialOrd for HuffmanNode {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

/// Build the Huffman tree from the static frequency table.
///
/// Uses a binary heap to repeatedly combine the two lowest-frequency
/// nodes until a single root remains.
#[must_use]
pub fn build_tree() -> HuffmanNode {
    let mut heap = BinaryHeap::new();

    for &(ch, freq) in FREQUENCIES {
        heap.push(HuffmanNode::Leaf { ch, freq });
    }

    while heap.len() > 1 {
        let left = heap.pop().unwrap();
        let right = heap.pop().unwrap();
        let freq = left.freq() + right.freq();
        heap.push(HuffmanNode::Internal {
            freq,
            left: Box::new(left),
            right: Box::new(right),
        });
    }

    heap.pop().unwrap()
}

/// Build a character-to-code mapping by walking the Huffman tree.
///
/// Each leaf is assigned a bit code where left edges append `0` and
/// right edges append `1`. Returns a map from character to a tuple
/// of `(bit_code, bit_length)`.
///
/// # Examples
///
/// ```rust
/// use srcdmp::raw::codec::huffman::HUFFMAN_CODES;
///
/// let codes = &*HUFFMAN_CODES;
/// // Common characters have short codes
/// if let Some(&(_code, length)) = codes.get(&' ') {
///     assert!(length < 8, "space should have a short Huffman code");
/// }
/// ```
#[must_use]
pub fn build_codes(node: &HuffmanNode) -> HashMap<char, (u64, u8)> {
    let mut codes = HashMap::new();
    walk(node, 0u64, 0u8, &mut codes);
    codes
}

fn walk(node: &HuffmanNode, code: u64, depth: u8, codes: &mut HashMap<char, (u64, u8)>) {
    match node {
        HuffmanNode::Leaf { ch, .. } => {
            codes.insert(*ch, (code, depth));
        }
        HuffmanNode::Internal { left, right, .. } => {
            walk(left, code << 1, depth + 1, codes);
            walk(right, (code << 1) | 1, depth + 1, codes);
        }
    }
}