// Length-limited canonical Huffman construction for the encode pipeline:
// frequencies in, (code lengths, bit-reversed code words) out.
//

///|
priv struct Node {
  lit : Int
  freq : Int
}

///|
priv struct LevelInfo {
  mut level : Int
  mut last_freq : Int
  mut next_char_freq : Int
  mut next_pair_freq : Int
  mut needed : Int
}

///|
fn LevelInfo::LevelInfo() -> LevelInfo {
  { level: 0, last_freq: 0, next_char_freq: 0, next_pair_freq: 0, needed: 0, }
}

///|
fn reverse_bits(number : UInt, bit_length : Int) -> UInt {
  reverse16(number << (16 - bit_length))
}

///|
/// Number of codes assigned to each bit length (index = length). Requires
/// `nodes` sorted by ascending frequency, with `n >= 3`.
fn bit_counts(
  nodes : Array[Node],
  n : Int,
  max_bits_in : Int,
) -> FixedArray[Int] {
  let maxi = 2147483647
  let mut max_bits = max_bits_in
  if max_bits > n - 1 {
    max_bits = n - 1
  }
  let levels = FixedArray::makei(17, _ => LevelInfo())
  let leaf_counts = FixedArray::make(17 * 17, 0)
  for level in 1..<=max_bits {
    levels[level].level = level
    levels[level].last_freq = nodes[1].freq
    levels[level].next_char_freq = nodes[2].freq
    levels[level].next_pair_freq = nodes[0].freq + nodes[1].freq
    levels[level].needed = 0
    leaf_counts[level * 17 + level] = 2
    if level == 1 {
      levels[level].next_pair_freq = maxi
    }
  }
  levels[max_bits].needed = 2 * n - 4
  let mut level = max_bits
  for ;; {
    let lv = levels[level]
    if lv.next_pair_freq == maxi && lv.next_char_freq == maxi {
      lv.needed = 0
      levels[level + 1].next_pair_freq = maxi
      level = level + 1
      continue
    }
    let prev_freq = lv.last_freq
    if lv.next_char_freq < lv.next_pair_freq {
      let nn = leaf_counts[level * 17 + level] + 1
      lv.last_freq = lv.next_char_freq
      leaf_counts[level * 17 + level] = nn
      lv.next_char_freq = if nn < n { nodes[nn].freq } else { maxi }
    } else {
      lv.last_freq = lv.next_pair_freq
      for i in 0.. 0 {
        level = level - 1
      }
    }
  }
  let bit_count = FixedArray::make(max_bits + 1, 0)
  let counts = leaf_counts[max_bits * 17:(max_bits + 1) * 17]
  let mut bits = 1
  for lvl = max_bits; lvl > 0; lvl = lvl - 1 {
    bit_count[bits] = counts[lvl] - counts[lvl - 1]
    bits = bits + 1
  }
  bit_count
}

///|
fn assign_codes(
  bit_count : FixedArray[Int],
  nodes : Array[Node],
  code_len : FixedArray[Int],
  code_val : FixedArray[UInt]?,
) -> Unit {
  // Assign lengths in frequency order, then generate canonical codes in
  // symbol order. No per-length arrays or sorting are needed.
  let mut rem = nodes.length()
  for nlen in 1.. 0 {
      values[symbol] = reverse_bits(next_code[nlen].reinterpret_as_uint(), nlen)
      next_code[nlen] += 1
    }
  }
}

///|
/// Build length-limited canonical Huffman codes for `size` symbols from their
/// frequencies into caller-owned arrays. Omit `code_val` when only lengths are
/// needed; this avoids both the unused code array and a tuple return allocation.
fn gen_huffman(
  freq : FixedArray[Int],
  size : Int,
  max_bits : Int,
  code_len : FixedArray[Int],
  code_val? : FixedArray[UInt],
) -> Unit {
  let nodes : Array[Node] = []
  for i in 0.. 2 else {
    for k in 0.. {
    if a.freq != b.freq {
      a.freq - b.freq
    } else {
      a.lit - b.lit
    }
  })
  let bit_count = bit_counts(nodes, count, max_bits)
  assign_codes(bit_count, nodes, code_len, code_val)
}