# Words by the Numbers 6, Exercise 3
import pathlib, re, urllib.request

URL = "https://www.gutenberg.org/cache/epub/1404/pg1404.txt"
DISPUTED = set(range(49, 59)) | {62, 63}   # the twelve disputed essays
JOINT = {18, 19, 20}                       # joint work, left out of the comparisons

def load(path="federalist.txt"):
    p = pathlib.Path(path)
    if not p.exists():                     # download once, then reuse the saved copy
        req = urllib.request.Request(URL, headers={"User-Agent": "Mozilla/5.0"})
        p.write_bytes(urllib.request.urlopen(req).read())
    raw = p.read_text(encoding="utf-8-sig").replace("\r", "")
    parts = re.split(r"^FEDERALIST No\. (\d+)\s*$", raw, flags=re.M)
    papers = {}                            # essay number -> (author, list of words)
    for k in range(1, len(parts), 2):
        n, body = int(parts[k]), parts[k + 1]
        m = re.search(r"^(HAMILTON|MADISON|JAY)$", body, re.M)
        author = m.group(1) if m else "?"
        if n in DISPUTED: author = "DISPUTED"
        if n in JOINT: author = "JOINT"
        start = body.find("To the People")
        end = body.rfind("\nPUBLIUS")
        text = body[start:end if end > start else None]
        papers[n] = (author, re.findall(r"[a-z]+(?:'[a-z]+)?", text.lower()))
    return papers

import statistics
from collections import Counter

papers = load()
known = {n: a for n, (a, w) in papers.items() if a in ("HAMILTON", "MADISON")}
counts = {n: Counter(w) for n, (a, w) in papers.items()}
size = {n: len(w) for n, (a, w) in papers.items()}

def make_model(train, k):
    pooled = Counter()
    for n in train:
        pooled.update(counts[n])
    vocab = [w for w, _ in pooled.most_common(k)]          # the k commonest words
    rates = [[counts[n][v] / size[n] for v in vocab] for n in train]
    mean = [statistics.mean(col) for col in zip(*rates)]
    sd = [statistics.pstdev(col) for col in zip(*rates)]
    profiles = {}
    for author in ("HAMILTON", "MADISON"):
        own = Counter()
        for n in train:
            if known[n] == author:
                own.update(counts[n])
        total = sum(own.values())
        profiles[author] = [(own[v] / total - mean[i]) / sd[i] for i, v in enumerate(vocab)]
    return vocab, mean, sd, profiles

def delta(n, model):
    vocab, mean, sd, profiles = model
    z = [(counts[n][v] / size[n] - mean[i]) / sd[i] for i, v in enumerate(vocab)]
    return {a: statistics.mean(abs(z[i] - p[i]) for i in range(len(vocab))) for a, p in profiles.items()}


for K in (10, 30, 100, 200):
    right = 0
    for n, a in known.items():
        model = make_model([m for m in known if m != n], K)
        d = delta(n, model)
        right += min(d, key=d.get) == a
    model = make_model(list(known), K)
    tally = Counter()
    for n in sorted(papers):
        if papers[n][0] == "DISPUTED":
            d = delta(n, model)
            tally[min(d, key=d.get)] += 1
    print(f"top {K:3d} words: {right} of {len(known)} known essays right; disputed essays -> {dict(tally)}")

# Output when run (Python 3.14; Federalist text from the Project Gutenberg plain-text file):
# top  10 words: 55 of 65 known essays right; disputed essays -> {'MADISON': 8, 'HAMILTON': 4}
# top  30 words: 61 of 65 known essays right; disputed essays -> {'MADISON': 11, 'HAMILTON': 1}
# top 100 words: 61 of 65 known essays right; disputed essays -> {'MADISON': 11, 'HAMILTON': 1}
# top 200 words: 61 of 65 known essays right; disputed essays -> {'MADISON': 12}