use std::io::Write;

#[derive(Debug, Clone)]
enum HT {
    Branch(Box<HT>, Box<HT>),
    Terminal(u8),
}

fn write_huffman(buf: &[u8], w: &mut impl Write) -> usize {
    let mut w = BitIO::new(w);

    let mut probs = [0usize; 256];
    for b in buf {
        probs[*b as usize] += 1;
    }

    let tree = HT::from_probabilities(probs);
    let mut table = [0u16; 256];
    tree.create_lut(&mut table, 1);

    for b in buf {
        let mut k = table[*b as usize];
        while k != 1 {
            w.wbit(k & 1 == 1);
            k <<= 1;
        }
    }

    w.position
}

impl HT {
    pub fn from_probabilities(ps: [usize; 256]) -> Self {
        let mut parts = ps
            .into_iter()
            .enumerate()
            .map(|(n, p)| (p, HT::Terminal(n as u8)))
            .collect::<Vec<_>>();

        while parts.len() != 1 {
            parts.sort_by_key(|e| -(e.0 as isize));
            let ((ap, at), (bp, bt)) = (parts.pop().unwrap(), parts.pop().unwrap());
            parts.push((ap + bp, HT::Branch(box at, box bt)))
        }
        parts[0].1.clone()
    }
    pub fn create_lut(&self, table: &mut [u16; 256], prefix: u16) {
        match self {
            HT::Branch(a, b) => {
                a.create_lut(table, (prefix << 1) | 0);
                b.create_lut(table, (prefix << 1) | 1);
            }
            HT::Terminal(n) => {
                assert_eq!(table[*n as usize], 0);
                table[*n as usize] = prefix
            }
        }
    }
}

struct BitIO<T> {
    inner: T,
    byte: u8,
    position: usize,
}
impl<T> BitIO<T> {
    pub fn new(inner: T) -> Self {
        Self {
            inner,
            byte: 0,
            position: 0,
        }
    }
}
impl<T: Write> BitIO<T> {
    #[inline]
    pub fn wbit(&mut self, b: bool) {
        self.byte <<= 1;
        self.byte |= b as u8;
        self.position += 1;
        if self.position & 0b111 == 0 {
            self.inner.write_all(&[self.byte]).unwrap();
        }
    }
}