# Words by the Numbers 9, Exercise 2
# Build a real Huffman code from the letter counts of Alice in Wonderland (cleaned to a-z
# and spaces), encode the first 20,000 characters as a string of 0s and 1s, decode it
# again, and compare the bits per character with the entropy and with a fixed 5-bit code.
import heapq, math, pathlib, re, urllib.request
from collections import Counter

def load(number):
    path = pathlib.Path(f"pg{number}.txt")
    if not path.exists():                       # download once, then reuse the saved copy
        url = f"https://www.gutenberg.org/cache/epub/{number}/pg{number}.txt"
        req = urllib.request.Request(url, headers={"User-Agent": "Mozilla/5.0"})
        path.write_bytes(urllib.request.urlopen(req).read())
    raw = path.read_text(encoding="utf-8-sig").replace("\r", "")
    start = raw.index("\n", raw.index("*** START OF")) + 1
    return raw[start:raw.index("*** END OF")]

text = re.sub(r"[^a-z]+", " ", load(11).lower()).strip()
counts = Counter(text)

heap = [(c, i, s) for i, (s, c) in enumerate(sorted(counts.items()))]   # leaves are (count, tie-break, symbol)
heapq.heapify(heap)
n = len(heap)
while len(heap) > 1:
    c1, _, a = heapq.heappop(heap)
    c2, _, b = heapq.heappop(heap)
    n += 1
    heapq.heappush(heap, (c1 + c2, n, (a, b)))   # a node is a pair of subtrees

code = {}
def walk(node, prefix=""):
    if isinstance(node, str):
        code[node] = prefix or "0"
    else:
        walk(node[0], prefix + "0")
        walk(node[1], prefix + "1")
walk(heap[0][2])

sample = text[:20000]
bits = "".join(code[ch] for ch in sample)
decode = {v: k for k, v in code.items()}
out, buffer = [], ""
for b in bits:                                   # no code is the start of another, so the first match is right
    buffer += b
    if buffer in decode:
        out.append(decode[buffer])
        buffer = ""
print("round trip ok:", "".join(out) == sample)
h = -sum(c / len(text) * math.log2(c / len(text)) for c in counts.values())
print(f"{len(sample):,} characters became {len(bits):,} bits = {len(bits) / len(sample):.3f} bits per character")
print(f"entropy of the whole book {h:.3f}; fixed-length code 5.000")
print("codes:", ", ".join(f"{'space' if s == ' ' else s} {code[s]}" for s in (" ", "e", "t", "a", "q", "z")))

# Output when run (Python 3.14; texts from the Project Gutenberg plain-text files):
# round trip ok: True
# 20,000 characters became 81,511 bits = 4.076 bits per character
# entropy of the whole book 4.046; fixed-length code 5.000
# codes: space 00, e 010, t 1101, a 1010, q 1100111011, z 1100111000