# Words by the Numbers 9, Exercise 1
# Surprise per letter. For Pride and Prejudice (cleaned to a-z and spaces):
# print the three least and three most surprising symbols, confirm that entropy is the
# frequency-weighted average of the surprises, and compare the surprise of "u" in general
# with its surprise straight after "q".
import 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(1342).lower()).strip()
counts = Counter(text)
total = len(text)
surprise = {s: -math.log2(c / total) for s, c in counts.items()}
ranked = sorted(surprise.items(), key=lambda kv: kv[1])

label = lambda s: "space" if s == " " else s
print("least surprising:", ", ".join(f"{label(s)} {b:.2f}" for s, b in ranked[:3]))
print("most surprising: ", ", ".join(f"{label(s)} {b:.2f}" for s, b in ranked[-3:]))
average = sum(counts[s] / total * surprise[s] for s in counts)
print(f"frequency-weighted average surprise = entropy = {average:.3f} bits")

after_q = Counter(text[i + 1] for i in range(total - 1) if text[i] == "q")
p_u = after_q["u"] / sum(after_q.values())
print(f"'u' anywhere: {surprise['u']:.2f} bits;  'u' right after 'q': {-math.log2(p_u):.2f} bits  ({after_q['u']} of {sum(after_q.values())} q's)")

# Output when run (Python 3.14; texts from the Project Gutenberg plain-text files):
# least surprising: space 2.43, e 3.25, t 3.81
# most surprising:  z 9.48, j 9.56, q 10.06
# frequency-weighted average surprise = entropy = 4.098 bits
# 'u' anywhere: 5.46 bits;  'u' right after 'q': 0.00 bits  (648 of 649 q's)