# Words by the Numbers 8, Exercise 2
import math, 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}

def load(number):
    path = pathlib.Path(f"pg{number}.txt")
    if not path.exists():
        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
    text = raw[start:raw.index("*** END OF")]
    return text.translate({0x2018: 39, 0x2019: 39, 0x201c: 34, 0x201d: 34})

def tokens(text):
    return re.findall(r"[a-z]+(?:'[a-z]+)?", text.lower())

SIZE = 500
pieces = []                                    # (book, token list) for every 500-word piece
for title, number in BOOKS.items():
    words = load(number).split()
    for i in range(0, len(words) - SIZE + 1, SIZE):
        pieces.append((title, tokens(" ".join(words[i:i + SIZE]))))
N = len(pieces)
df = Counter(w for _, ws in pieces for w in set(ws))

def keywords(title, max_share, top=8):
    total = defaultdict(float)
    for book, ws in pieces:
        if book != title:
            continue
        for w, k in Counter(ws).items():
            if df[w] <= max_share * N:
                total[w] += (k / len(ws)) * math.log(N / df[w])
    return [w for w, _ in sorted(total.items(), key=lambda kv: -kv[1])[:top]]

for title in BOOKS:
    print(title)
    for max_share in (1.0, 0.2):
        print(f"  words allowed in at most {max_share:.0%} of pieces: " + ", ".join(keywords(title, max_share)))

# Output when run (Python 3.14; texts from the Project Gutenberg plain-text files):
# Alice in Wonderland
#   words allowed in at most 100% of pieces: alice, she, said, you, turtle, mock, queen, gryphon
#   words allowed in at most 20% of pieces: alice, turtle, mock, queen, gryphon, hatter, mouse, rabbit
# The Time Machine
#   words allowed in at most 100% of pieces: i, my, me, machine, traveller, was, morlocks, came
#   words allowed in at most 20% of pieces: machine, traveller, morlocks, came, weena, white, saw, seemed
# Pride and Prejudice
#   words allowed in at most 100% of pieces: her, she, you, he, mr, i, his, elizabeth
#   words allowed in at most 20% of pieces: bennet, bingley, miss, jane, lady, collins, wickham, illustration
# The Federalist Papers
#   words allowed in at most 100% of pieces: states, government, state, power, executive, constitution, national, people
#   words allowed in at most 20% of pieces: executive, powers, senate, legislative, representatives, members, laws, number