# Words by the Numbers 9, Exercise 3
# Play Shannon's guessing game on a 60-character passage of Alice in Wonderland with a
# machine that has only read the other three books. For context lengths 0 to 3 print the
# average number of guesses over the whole book, then for order 2 print the number of
# guesses needed for each character of the passage (a long run of 1s means easy text).
import pathlib, re, urllib.request
from collections import Counter, defaultdict

BOOKS = {"Alice in Wonderland": 11, "The Time Machine": 35,
         "Pride and Prejudice": 1342, "The Federalist Papers": 1404}
ALPHABET = " abcdefghijklmnopqrstuvwxyz"

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")]

texts = {t: re.sub(r"[^a-z]+", " ", load(n).lower()).strip() for t, n in BOOKS.items()}
pool = "".join(t for title, t in texts.items() if title != "Alice in Wonderland")
test = texts["Alice in Wonderland"]
overall = Counter(pool)

def guesses(k):
    model = defaultdict(Counter)
    for i in range(k, len(pool)):
        model[pool[i - k:i]][pool[i]] += 1
    out = []
    for i in range(k, len(test)):
        c = model.get(test[i - k:i], Counter())
        order = sorted(ALPHABET, key=lambda ch: (-c[ch], -overall[ch]))
        out.append(order.index(test[i]) + 1)
    return out

for k in range(4):
    g = guesses(k)
    print(f"order {k}: average {sum(g) / len(g):.2f} guesses, first guess right {g.count(1) / len(g):.1%}")

g = guesses(2)                                   # g[j] belongs to test[j + 2]
start = 5000
passage = test[start:start + 60]
print(repr(passage))
print(" ".join(str(x) for x in g[start - 2:start - 2 + 60]))

# Output when run (Python 3.14; texts from the Project Gutenberg plain-text files):
# order 0: average 7.00 guesses, first guess right 20.2%
# order 1: average 4.75 guesses, first guess right 29.1%
# order 2: average 3.73 guesses, first guess right 39.3%
# order 3: average 3.31 guesses, first guess right 47.1%
# 'either question it didn t much matter which way she put it s'
# 14 9 4 2 1 2 1 25 1 3 3 3 3 1 1 1 4 2 1 15 2 6 12 2 1 9 9 5 2 1 1 8 1 3 4 1 1 1 6 1 1 2 1 1 5 3 3 1 6 1 1 1 7 6 3 1 2 2 1 8