# Words by the Numbers 8, Exercise 1
import math
from collections import Counter

docs = {                                   # invented, illustrative mini-documents
    "A": "the cat sat on the mat and the cat purred",
    "B": "the dog sat on the log and the dog barked",
    "C": "the bank raised the interest rate on the loan",
    "D": "the cat chased the dog around the bank",
}
words = {name: text.split() for name, text in docs.items()}
N = len(docs)
df = Counter(w for ws in words.values() for w in set(ws))

variants = {
    "plain  ln(N/df)": lambda d: math.log(N / d),
    "smooth ln(1+N/df)": lambda d: math.log(1 + N / d),
}
print("idf for each document frequency (N = 4):")
for label, f in variants.items():
    print(f"  {label:18}", "  ".join(f"df={d}: {f(d):.3f}" for d in range(1, N + 1)))

for label, f in variants.items():
    print(f"top 3 per document, {label}")
    for name, ws in words.items():
        counts = Counter(ws)
        score = {w: (c / len(ws)) * f(df[w]) for w, c in counts.items()}
        top = sorted(score.items(), key=lambda kv: (-kv[1], kv[0]))[:3]
        print(f"  {name}: " + ", ".join(f"{w} {s:.3f}" for w, s in top))

# Output when run (Python 3.14; texts from the Project Gutenberg plain-text files):
# idf for each document frequency (N = 4):
#   plain  ln(N/df)    df=1: 1.386  df=2: 0.693  df=3: 0.288  df=4: 0.000
#   smooth ln(1+N/df)  df=1: 1.609  df=2: 1.099  df=3: 0.847  df=4: 0.693
# top 3 per document, plain  ln(N/df)
#   A: cat 0.139, mat 0.139, purred 0.139
#   B: barked 0.139, dog 0.139, log 0.139
#   C: interest 0.154, loan 0.154, raised 0.154
#   D: around 0.173, chased 0.173, bank 0.087
# top 3 per document, smooth ln(1+N/df)
#   A: cat 0.220, the 0.208, mat 0.161
#   B: dog 0.220, the 0.208, barked 0.161
#   C: the 0.231, interest 0.179, loan 0.179
#   D: the 0.260, around 0.201, chased 0.201