///|
priv struct BitOutput {
  sink : Sink
  mut carry : UInt64
  mut pending : Int32
  mut count : Int32
}

///|
fn bit_output() -> BitOutput {
  { sink: sink(), carry: 0UL, pending: 0, count: 0, }
}

///|
fn BitOutput::put(self : BitOutput, value : UInt64, width : Int32) -> Unit {
  self.count += width
  self.carry = (self.carry << width) | value
  self.pending += width
  while self.pending >= 8 {
    self.pending -= 8
    self.sink.bytes.push((self.carry >> self.pending).to_byte())
  }
}

///|
fn BitOutput::finish(self : BitOutput) -> Bytes {
  if self.pending > 0 {
    self.sink.bytes.push((self.carry << (8 - self.pending)).to_byte())
  }
  self.sink.finish()
}

///|
priv struct BitInput {
  bytes : Bytes
  limit : Int32
  mut pos : Int32
}

///|
fn BitInput::get(self : BitInput, width : Int32) -> UInt64 raise ExrError {
  if width > self.limit - self.pos {
    raise Invalid("truncated PIZ Huffman bits", self.pos / 8)
  }
  let mut n = 0UL
  let end = self.pos + width
  while self.pos < end {
    let shift = 8 - self.pos % 8
    let count = imin(shift, end - self.pos)
    let bits = self.bytes[self.pos / 8].to_uint64() >> (shift - count)
    n = (n << count) | (bits & ((1UL << count) - 1UL))
    self.pos += count
  }
  n
}

///|
priv struct Canonical {
  counts : Array[Int32]
  starts : Array[UInt64]
  offsets : Array[Int32]
  symbols : Array[Int32]
  codes : Array[UInt64]
  longest : Int32
}

///|
fn canonical(lengths : Array[Int32]) -> Canonical raise ExrError {
  let counts = Array::make(59, 0)
  let mut longest = 0
  for size in lengths {
    if size > 0 {
      counts[size] += 1
      longest = imax(longest, size)
    }
  }
  let starts = Array::make(59, 0UL)
  let offsets = Array::make(59, 0)
  let mut next = 0UL
  for size = longest; size > 0; size = size - 1 {
    starts[size] = next
    let end = next + counts[size].to_uint64()
    if end > 1UL << size || (end & 1UL) != 0UL {
      raise Invalid("invalid PIZ Huffman prefix lengths", 0)
    }
    next = end >> 1
  }
  if next != 1UL {
    raise Invalid("incomplete PIZ Huffman table", 0)
  }
  for size = 2; size <= longest; size = size + 1 {
    offsets[size] = offsets[size - 1] + counts[size - 1]
  }
  let used = Array::make(59, 0)
  let symbols = Array::make(offsets[longest] + counts[longest], 0)
  let codes = Array::make(lengths.length(), 0UL)
  for i = 0; i < lengths.length(); i = i + 1 {
    let size = lengths[i]
    if size > 0 {
      codes[i] = starts[size] + used[size].to_uint64()
      symbols[offsets[size] + used[size]] = i
      used[size] += 1
    }
  }
  { counts, starts, offsets, symbols, codes, longest, }
}

///|
fn BitInput::symbol(self : BitInput, table : Canonical) -> Int32 raise ExrError {
  let mut code = 0UL
  for size = 1; size <= table.longest; size = size + 1 {
    code = (code << 1) | self.get(1)
    let first = table.starts[size]
    if code >= first && code - first < table.counts[size].to_uint64() {
      return table.symbols[table.offsets[size] + (code - first).to_int()]
    }
  }
  raise Invalid("invalid PIZ Huffman code", self.pos / 8)
}

///|
fn huffman_decode(bytes : Bytes, count : Int32) -> Array[Int32] raise ExrError {
  let r = reader(bytes)
  let first = r.i32()
  let run = r.i32()
  let table_bytes = r.i32()
  let data_bits = r.i32()
  let reserved = r.i32()
  if first < 0 ||
    first >= run ||
    run > 65536 ||
    table_bytes < 1 ||
    table_bytes > 49153 ||
    data_bits < 1 ||
    reserved != 0 {
    raise Invalid("PIZ Huffman header", 0)
  }
  let lengths = Array::make(run - first + 1, 0)
  let bits = { bytes: r.take(table_bytes), limit: table_bytes * 8, pos: 0, }
  let mut i = 0
  while i < lengths.length() {
    let size = bits.get(6).to_int()
    if size >= 59 {
      let zeros = if size == 63 { bits.get(8).to_int() + 6 } else { size - 57 }
      if zeros > lengths.length() - i {
        raise Invalid("PIZ Huffman table zero run", 0)
      }
      i += zeros
    } else {
      lengths[i] = size
      i += 1
    }
  }
  if lengths[0] == 0 || lengths[run - first] == 0 || bits.limit - bits.pos >= 8 {
    raise Invalid("PIZ Huffman table length", 0)
  }
  let table = canonical(lengths)
  let data = {
    bytes: r.take((data_bits - 1) / 8 + 1),
    limit: data_bits,
    pos: 0,
  }
  if r.pos != r.end {
    raise Invalid("PIZ Huffman trailing bytes", r.pos)
  }
  let result = Array::make(count, 0)
  let mut written = 0
  while data.pos < data.limit {
    let symbol = first + data.symbol(table)
    if symbol == run {
      let repeats = data.get(8).to_int()
      if written == 0 || repeats > count - written {
        raise Invalid("PIZ Huffman repeat exceeds sample count", 0)
      }
      let prior = result[written - 1]
      for j = 0; j < repeats; j = j + 1 {
        result[written + j] = prior
      }
      written += repeats
    } else {
      if written == count {
        raise Invalid("PIZ Huffman excess samples", 0)
      }
      result[written] = symbol
      written += 1
    }
  }
  if written != count {
    raise Invalid("PIZ Huffman short sample count", 0)
  }
  result
}

///|
/// Sort the leaves once and merge the two lightest nodes at each step.
fn huffman_lengths(
  words : Array[Int32],
  first : Int32,
  run : Int32,
) -> Array[Int32] {
  let leaves = run - first + 1
  let freq = Array::make(leaves, 0)
  for word in words {
    freq[word - first] += 1
  }
  freq[leaves - 1] = 1
  let sorted = []
  for i = 0; i < leaves; i = i + 1 {
    if freq[i] > 0 {
      sorted.push(i)
    }
  }
  sorted.sort_by((a, b) => {
    let order = freq[a].compare(freq[b])
    if order == 0 {
      a.compare(b)
    } else {
      order
    }
  })
  let nodes = leaves + sorted.length() - 1
  let weight = Array::make(nodes, 0)
  for i in sorted {
    weight[i] = freq[i]
  }
  let depth = Array::make(nodes, 0)
  let mut leaf = 0
  let mut branch = leaves
  for parent = leaves; parent < nodes; parent = parent + 1 {
    for _child = 0; _child < 2; _child = _child + 1 {
      let take_leaf = leaf < sorted.length() &&
        (branch == parent || weight[sorted[leaf]] <= weight[branch])
      let node = if take_leaf {
        let index = sorted[leaf]
        leaf += 1
        index
      } else {
        let index = branch
        branch += 1
        index
      }
      weight[parent] += weight[node]
      depth[node] = parent
    }
  }
  for i = nodes - 2; i >= 0; i = i - 1 {
    if depth[i] != 0 {
      depth[i] = depth[depth[i]] + 1
    }
  }
  depth[:leaves].to_owned()
}

///|
fn huffman_encode(
  words : Array[Int32],
  budget : Int32,
) -> Bytes? raise ExrError {
  let mut first = 65535
  let mut largest = 0
  for word in words {
    first = imin(first, word)
    largest = imax(largest, word)
  }
  let run = largest + 1
  let lengths = huffman_lengths(words, first, run)
  let table = canonical(lengths)
  let descriptor = bit_output()
  let mut i = 0
  while i < lengths.length() {
    if lengths[i] == 0 {
      let mut n = 1
      while n < 261 && i + n < lengths.length() && lengths[i + n] == 0 {
        n += 1
      }
      if n >= 6 {
        descriptor.put(63UL, 6)
        descriptor.put((n - 6).to_uint64(), 8)
      } else {
        descriptor.put((if n == 1 { 0 } else { n + 57 }).to_uint64(), 6)
      }
      i += n
    } else {
      descriptor.put(lengths[i].to_uint64(), 6)
      i += 1
    }
  }
  let encoded_table = descriptor.finish()
  let available = budget - 20 - encoded_table.length()
  if available <= 0 {
    return None
  }
  let bit_budget = imin(2147483647 / 8, available) * 8
  let data = bit_output()
  let run_size = lengths[run - first]
  i = 0
  while i < words.length() {
    let symbol = words[i] - first
    let size = lengths[symbol]
    let mut n = 1
    while n < 256 && i + n < words.length() && words[i + n] == words[i] {
      n += 1
    }
    let repeat = run_size + 8 < size * (n - 1)
    let cost = if repeat { size + run_size + 8 } else { size * n }
    if cost > bit_budget - data.count {
      return None
    }
    if repeat {
      data.put(table.codes[symbol], size)
      data.put(table.codes[run - first], run_size)
      data.put((n - 1).to_uint64(), 8)
    } else {
      for j = 0; j < n; j = j + 1 {
        data.put(table.codes[symbol], size)
      }
    }
    i += n
  }
  let out = sink()
  out.i32(first)
  out.i32(run)
  out.i32(encoded_table.length())
  out.i32(data.count)
  out.i32(0)
  out.raw(encoded_table)
  out.raw(data.finish())
  Some(out.finish())
}