# Words by the Numbers 8, Exercise 3
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
docs = []                                      # (book, piece number)
index = defaultdict(dict)                      # word -> {doc id: count}
for title, number in BOOKS.items():
    words = load(number).split()
    for i in range(0, len(words) - SIZE + 1, SIZE):
        for w, k in Counter(tokens(" ".join(words[i:i + SIZE]))).items():
            index[w][len(docs)] = k
        docs.append((title, i // SIZE))
N = len(docs)

SCHEMES = {                                    # (count weight, use idf?)
    "raw count":        (lambda k: k, False),
    "log count":        (lambda k: 1 + math.log(k), False),
    "log count x idf":  (lambda k: 1 + math.log(k), True),
}

def search(query, scheme, top=10):
    weight, use_idf = SCHEMES[scheme]
    scores = defaultdict(float)
    for w in tokens(query):
        if w not in index:
            continue
        idf = math.log(N / len(index[w])) if use_idf else 1.0
        for doc, k in index[w].items():
            scores[doc] += weight(k) * idf
    return sorted(scores.items(), key=lambda kv: -kv[1])[:top]

QUERIES = [                                    # (query, book I expect it to come from)
    ("what did the rabbit say to the queen", "Alice in Wonderland"),
    ("the cat and the hatter", "Alice in Wonderland"),
    ("the traveller and the time machine", "The Time Machine"),
    ("the weena and the morlocks", "The Time Machine"),
    ("mr darcy and the ball at netherfield", "Pride and Prejudice"),
    ("elizabeth and her sister jane", "Pride and Prejudice"),
    ("the powers of the national government", "The Federalist Papers"),
    ("the people and the states of the union", "The Federalist Papers"),
]
print(f"{'query':42}" + "".join(f"{s:>17}" for s in SCHEMES) + "   (of top 10 from the expected book)")
totals = Counter()
for query, book in QUERIES:
    row = f"{query:42}"
    for scheme in SCHEMES:
        right = sum(docs[doc][0] == book for doc, _ in search(query, scheme))
        totals[scheme] += right
        row += f"{right:17d}"
    print(row)
print(f"{'total out of ' + str(10 * len(QUERIES)):42}" + "".join(f"{totals[s]:17d}" for s in SCHEMES))

# Output when run (Python 3.14; texts from the Project Gutenberg plain-text files):
# query                                             raw count        log count  log count x idf   (of top 10 from the expected book)
# what did the rabbit say to the queen                      0               10               10
# the cat and the hatter                                    0                9               10
# the traveller and the time machine                        0               10               10
# the weena and the morlocks                                0               10               10
# mr darcy and the ball at netherfield                      0               10               10
# elizabeth and her sister jane                             9               10               10
# the powers of the national government                    10               10               10
# the people and the states of the union                   10               10               10
# total out of 80                                          29               79               80