hfmn 0.1.0

use crate::code::HfmnCode;
use crate::errors::Error;
use crate::table::DecodeTables;
use crate::tree::{get_smallest_node, Node};
use itertools::Itertools;
use rustc_hash::FxHashMap as HashMap;
use serde::{Deserialize, Serialize};
use std::fmt::{self, Debug, Display};
use std::{collections::VecDeque, hash::Hash};

/// A mapping of huffman codes to data symbols.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct CodeBook<T: Hash + Eq> {
    /// The forward mapping used by the encoder
    coding: HashMap<T, HfmnCode>,
    /// The stack of decoding tables
    decoding: DecodeTables<T>,
}

impl<T: Hash + Eq + Clone + Ord + Debug> CodeBook<T> {
    /// Constructs a code book optimal over the given symbols.
    ///
    /// The codebook will provide the optimal single-symbol prefix-free mapping,
    /// and can be used to encode or decode data according to that mapping.
    ///
    /// # Arguments
    /// * `data` - A slice of symbols.
    ///
    /// # Errors
    /// * `Error::DataEmpty` if `data` contains 0 or 1 distinct symbols, since
    ///   there is no sense in encoding this
    pub fn from_data(data: impl AsRef<[T]>) -> Result<Self, Error> {
        let mut counts: HashMap<T, usize> = HashMap::default();

        for d in data.as_ref() {
            counts
                .entry(d.clone())
                .and_modify(|c| *c = c.saturating_add(1))
                .or_insert(1);
        }

        let mut leaf_nodes = counts
            .into_iter()
            .sorted_by(|(_, a), (_, b)| a.cmp(b))
            .sorted_by(|(a, _), (b, _)| a.cmp(b))
            .map(|(a, b)| {
                #[allow(clippy::cast_precision_loss)]
                // The precision loss here is fine since we are just using it for a probability, if
                // the results are really that close then there's no real-world difference
                {
                    Node::new(None, (b as f64) / (data.as_ref().len() as f64), Some(a))
                }
            })
            .collect::<VecDeque<Node<T>>>();

        let mut internal_nodes = VecDeque::new();

        while leaf_nodes.len().saturating_add(internal_nodes.len()) >= 2 {
            let node_1 = get_smallest_node(&mut leaf_nodes, &mut internal_nodes);
            let node_2 = get_smallest_node(&mut leaf_nodes, &mut internal_nodes);
            let combined_probability = node_1.probability() + node_2.probability();

            internal_nodes.push_back(Node::new(
                Some((node_1, node_2)),
                combined_probability,
                None,
            ));
        }

        let Some(root_node) = internal_nodes.pop_front() else {
            // We should always have exactly one node left in the internal
            // queue, so we only reach this if no nodes were added to start with
            return Err(Error::DataEmpty);
        };

        let mut map = HashMap::default();
        root_node.traverse(&mut map, HfmnCode::new());

        Ok(Self {
            decoding: DecodeTables::new(&map),
            coding: map,
        })
    }

    /// Returns a codebook crated from a supplied mapping.
    ///
    /// # Arguments
    /// * `mapping` - A slice of (symbol, code). This must contain each symbol
    ///   at most once, and each ocde must be valid, and the set of codes must
    ///   be correctly formed.
    pub fn from_mapping(mapping: impl AsRef<[(T, HfmnCode)]>) -> Self {
        let coding = mapping
            .as_ref()
            .iter()
            .cloned()
            .collect::<HashMap<T, HfmnCode>>();

        Self {
            decoding: DecodeTables::new(&coding),
            coding,
        }
    }

    /// Encodes symbols using the calculated Huffman codebook
    ///
    /// Returns both the encoded version of the symbols and the total number of
    /// symbols encoded. This is needed by the decoder in order to accurately
    /// decode the codestream.
    ///
    /// # Arguments
    /// * `data` - The symbols to be encoded
    ///
    /// # Errors
    /// * `Error::SymbolNotFound` if any of the symbols in data are not in the
    ///   codebook
    pub fn encode_data(&self, data: impl AsRef<[T]>) -> Result<Vec<&HfmnCode>, Error> {
        data.as_ref()
            .iter()
            .map(|d| self.coding.get(d).ok_or(Error::SymbolNotFound))
            .collect::<Result<_, _>>()
    }

    /// Decode huffman codes
    ///
    /// # Arguments
    /// * `bytes` - The bytes containing the codestream
    /// * `num_symbols` - The number of symbols that were encoded
    ///
    /// # Errors
    /// * `Error::InvalidCode` if any of the codes read are not found in the
    ///   codebook
    pub fn decode_data(&self, bytes: &[u8], num_symbols: usize) -> Result<Vec<&T>, Error> {
        self.decoding.decode(bytes, num_symbols)
    }

    // pub fn from_codebook(symbols_lengths: impl AsRef<[(T, u8)]>) -> Result<Self,
    // Error> {     let mut mapping = HashMap::default();

    //     let mut current_symbol = BitVec::new();
    //     for (symbol, length) in symbols_lengths.as_ref().iter() {
    //         if *length as usize != current_symbol.len() {
    //             current_symbol.resize(*length as usize, false);
    //         }

    //         while
    //     }

    //     Ok()
    // }
}

#[derive(Debug, Serialize, Deserialize)]
struct InternalCoding<T> {
    pub codes: Vec<(T, u8)>,
}

impl<'a, T: Serialize + Deserialize<'a> + Hash + Ord + Clone + Debug> CodeBook<T> {
    #[must_use]
    /// Serialize the codebook to bytes
    pub fn encode_book(self) -> Vec<u8> {
        let internal_code = InternalCoding {
            codes: self
                .coding
                .into_iter()
                .map(|(symbol, code)| (symbol, code.len()))
                .sorted_by(|(a, _), (b, _)| a.cmp(b))
                .sorted_by(|(_, a), (_, b)| a.cmp(b))
                .collect(),
        };

        // TODO: Check this is always valid
        bincode::serialize(&internal_code).unwrap_or_else(|_| {
                unreachable!(
                    "Something was wrong with the internal serialization, please report a bug"
                )
            })
    }

    /// Deserialize a codebook from bytes
    ///
    /// Note this function is not intended to provide deep validation on whether
    /// the codebook is valid or not.
    ///
    /// # Arguments
    /// * `bytes` - The bytes to deserialize back into a codebook
    ///
    /// # Errors
    /// * `Error::BinaryCoding` - If the bytes can't be decoded into a trivially
    ///   valid codebook
    pub fn decode_book(bytes: &'a [u8]) -> Result<Self, Error> {
        let internal_code: InternalCoding<T> = bincode::deserialize(bytes)?;

        let mut coding = HashMap::default();
        let mut current_code = HfmnCode::new();
        for (symbol, length) in internal_code.codes {
            if length > current_code.len() {
                *current_code.inner_mut() <<= length - current_code.len();
                *current_code.length_mut() = length;
            }

            coding.insert(symbol, current_code);

            *current_code.inner_mut() += 1;
        }

        Ok(Self {
            decoding: DecodeTables::new(&coding),
            coding,
        })
    }
}

impl<T: Debug + Hash + Eq> Display for CodeBook<T> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "Huffman Coding: ")?;
        for (symbol, coding) in &self.coding {
            write!(f, "\n({symbol:?} -> {coding})")?;
        }
        Ok(())
    }
}

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

    #[test]
    fn empty_creation() {
        let data: [char; 0] = [];

        assert!(matches!(CodeBook::from_data(data), Err(Error::DataEmpty)));
    }

    #[test]
    fn empty_encode() {
        let data = ['a', 'b'];

        let codebook = CodeBook::from_data(data).unwrap();

        assert_eq!(codebook.encode_data([]).unwrap(), Vec::<&HfmnCode>::new());
    }

    #[test]
    fn single_symbol_encode() {
        let data = ['a', 'b'];

        let codebook = CodeBook::from_data(data).unwrap();

        assert_eq!(codebook.encode_data(['a']).unwrap().len(), 1);
    }

    #[test]
    fn invalid_symbol() {
        let data = ['a', 'b'];

        let codebook = CodeBook::from_data(data).unwrap();

        assert!(matches!(
            codebook.encode_data(['c']),
            Err(Error::SymbolNotFound)
        ));
    }

    #[test]
    fn duplicated_symbols() {
        let data = ['a', 'a'];

        assert!(matches!(CodeBook::from_data(data), Err(Error::DataEmpty)));
    }

    #[test]
    fn through_book_code() {
        let data = ['a', 'b', 'b'];

        let book = CodeBook::from_data(data).unwrap();
        let encoded_book = book.clone().encode_book();
        let decoded_book = CodeBook::<char>::decode_book(&encoded_book).unwrap();

        assert_eq!(book, decoded_book);
    }

    #[test]
    fn through_book_code_with_gap() {
        let data = ['a', 'c', 'c'];

        let book = CodeBook::from_data(data).unwrap();
        let encoded_book = book.clone().encode_book();
        let decoded_book = CodeBook::<char>::decode_book(&encoded_book).unwrap();

        assert_eq!(book, decoded_book);
    }
}