"""Build, encode and decode Huffman codes; compare average length with entropy."""
from __future__ import annotations
from collections import Counter
import heapq
import itertools
import math


def entropy(probabilities):
    """Pre: finite nonnegative probabilities summing to 1. Post: entropy in bits."""
    p = list(probabilities)
    if not p or any(not math.isfinite(x) or x < 0 for x in p) or not math.isclose(sum(p), 1):
        raise ValueError("probabilities must be nonnegative, finite, and sum to one")
    return -sum(value*math.log2(value) for value in p if value)


def huffman(counts):
    """Pre: character keys, positive integer counts. Post: prefix code; singleton uses '0'."""
    if (not counts or any(not isinstance(k, str) or len(k) != 1 for k in counts)
            or any(not isinstance(v, int) or v <= 0 for v in counts.values())):
        raise ValueError("counts need single-character keys and positive integers")
    serial = itertools.count()
    heap = [(weight, next(serial), symbol) for symbol, weight in sorted(counts.items())]
    heapq.heapify(heap)
    # Invariant: heap nodes partition symbols and carry the sum of their counts.
    # Bound: number of heap nodes minus one.
    while len(heap) > 1:
        a, _, left = heapq.heappop(heap)
        b, _, right = heapq.heappop(heap)
        heapq.heappush(heap, (a+b, next(serial), (left, right)))
    codes, pending = {}, [(heap[0][2], "")]
    # Invariant: processed leaves have their tree paths; pending nodes complete the tree.
    # Bound: remaining tree nodes.
    while pending:
        node, path = pending.pop()
        if isinstance(node, str):
            codes[node] = path or "0"
        else:
            pending.extend(((node[0], path+"0"), (node[1], path+"1")))
    return codes


def encode(text, codes):
    """Pre: every character has a code. Post: concatenated bit string."""
    return "".join(codes[character] for character in text)


def decode(bits, codes):
    """Pre: valid prefix map. Post: original text, or ValueError for malformed bits."""
    if set(bits) - {"0", "1"}:
        raise ValueError("encoded input must be binary")
    reverse, output, prefix = {word: symbol for symbol, word in codes.items()}, [], ""
    # Invariant: output plus prefix covers exactly the consumed bits.
    # Bound: remaining bits.
    for bit in bits:
        prefix += bit
        if prefix in reverse:
            output.append(reverse[prefix])
            prefix = ""
        elif not any(word.startswith(prefix) for word in reverse):
            raise ValueError("invalid code prefix")
    if prefix:
        raise ValueError("truncated codeword")
    return "".join(output)


def demo():
    text = "AAAAAAAABBBBCCDD"
    counts = Counter(text)
    codes = huffman(counts)
    bits = encode(text, codes)
    print("Codes:", codes)
    print("Encoded:", bits, "length:", len(bits))
    print("Decoded:", decode(bits, codes))
    print("Entropy:", entropy([count/len(text) for count in counts.values()]), "bits/symbol")
    print("Average code length:", len(bits)/len(text), "bits/symbol")
    print("Model and message framing costs are not included in the payload length.")


def self_test():
    count = 0
    for text in ("a", "aaaa", "ab", "abracadabra", "AAAAAAAABBBBCCDD", "banana"):
        codes = huffman(Counter(text))
        check(decode(encode(text, codes), codes), text)
        check(all(not b.startswith(a) for a in codes.values() for b in codes.values() if a != b), True)
        count += 2
    check(entropy([.5, .25, .125, .125]), 1.75)
    check(entropy([1, 0]), 0.0)
    return count + 2

def check(actual, expected):
    """Require an expected result; this check also runs under python -O."""
    if actual != expected:
        raise AssertionError(f"expected {expected!r}, got {actual!r}")


if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--self-test", action="store_true")
    args = parser.parse_args()
    if args.self_test:
        print(f"{self_test()} checks passed")
    else:
        demo()