"""tf-idf and PMI, worked out as Chapters 4 and 7 work them out.

Two ideas, both of which come down to dividing by what you expected.

    tf-idf  divides a word's count in a document by how many documents hold it
    PMI     divides a pair's joint count by what independence would predict

    python3 weighting.py tfidf     # the moon / the worked example
    python3 weighting.py corpus    # the four-document matrix and cosines
    python3 weighting.py pmi       # the ice cream worked example
    python3 weighting.py all       # all three

Run `pmi` and watch the word "the" fail to score, which is the whole point.
Run `corpus` and watch two documents that share only "the" come out at a
cosine similarity of exactly zero.

Install:  nothing, the Python standard library is enough
"""

import argparse
import itertools
import math
from collections import Counter


# --------------------------------------------------------------- tf-idf

def idf(n_docs, doc_freq, base=10):
    """How informative a term is across the collection.

    A term in every document scores 0, because N/df = 1 and log 1 = 0. It
    cannot tell any two documents apart, so its weight is annihilated.
    """
    return math.log(n_docs / doc_freq, base)


def tf_idf(count_in_doc, doc_length, n_docs, doc_freq, base=10):
    """Length-normalised term frequency times inverse document frequency."""
    tf = count_in_doc / doc_length
    return tf * idf(n_docs, doc_freq, base)


def demo_tfidf():
    N = 100_000
    print(f"\ntf-idf   collection of N = {N:,} documents")
    print("         one document, 427 tokens long\n")
    rows = [
        # word,      count in doc, document frequency
        ("moon",     20,   100),
        ("the",      40,   N),
        ("telescope", 3,   4_000),
        ("crater",    1,   60),
    ]
    print(f"  {'word':<12}{'count':>7}{'df':>9}{'tf':>9}{'idf':>7}{'tf-idf':>9}")
    print("  " + "-" * 53)
    for word, count, df in rows:
        tf = count / 427
        i = idf(N, df)
        print(f"  {word:<12}{count:>7}{df:>9,}{tf:>9.4f}{i:>7.2f}"
              f"{tf * i:>9.4f}")

    print("\n  Read the last two rows against the first two.")
    print("  'the' occurs twice as often as 'moon' and scores ZERO, because")
    print("  it is in every document. 'crater' occurs once and outscores it.")
    print("  Raw frequency is the wrong guide; tf-idf is a machine for")
    print("  discounting the Zipfian head.")


# -------------------------------------------------- the four-document corpus

DOCS = {
    "d1": "the moon orbits the earth",
    "d2": "the earth orbits the sun",
    "d3": "the telescope shows the moon",
    "d4": "the sun is a star",
}
# Listed in this order so the printed matrix matches the book's table.
VOCAB = ["the", "moon", "orbits", "earth", "sun",
         "telescope", "shows", "is", "a", "star"]


def cosine(u, v):
    """Angle between two vectors, ignoring their lengths."""
    nu = math.sqrt(sum(x * x for x in u))
    nv = math.sqrt(sum(x * x for x in v))
    if nu == 0 or nv == 0:
        return 0.0
    return sum(a * b for a, b in zip(u, v)) / (nu * nv)


def demo_corpus():
    names = list(DOCS)
    counts = {d: Counter(DOCS[d].split()) for d in names}
    n_docs = len(names)
    df = {w: sum(1 for d in names if counts[d][w] > 0) for w in VOCAB}
    idf = {w: math.log10(n_docs / df[w]) for w in VOCAB}

    print(f"\nA four-document collection, N = {n_docs}\n")
    for d in names:
        print(f"  {d}: {DOCS[d]}")

    print("\nStage 1. Count. Rows are terms, columns are documents.\n")
    head = "".join(f"{d:>6}" for d in names)
    print(f"  {'term':<11}{head}")
    print("  " + "-" * (11 + 6 * n_docs))
    for w in VOCAB:
        print(f"  {w:<11}" + "".join(f"{counts[d][w]:>6}" for d in names))

    print("\nStage 2. Count the documents, not the occurrences.")
    print("         df is how many documents hold the term at all.")
    print("         idf = log10(N / df).\n")
    print(f"  {'term':<11}{'df':>4}{'N/df':>8}{'idf':>8}")
    print("  " + "-" * 31)
    for w in VOCAB:
        print(f"  {w:<11}{df[w]:>4}{n_docs/df[w]:>8.2f}{idf[w]:>8.3f}")
    print("\n  'the' is in every document, so N/df = 1 and idf = 0.")
    print("  Its weight is annihilated before any similarity is computed.")

    print("\nStage 3. Multiply, cell by cell. tf-idf = tf x idf.\n")
    weights = {d: [counts[d][w] * idf[w] for w in VOCAB] for d in names}
    print(f"  {'term':<11}" + "".join(f"{d:>8}" for d in names))
    print("  " + "-" * (11 + 8 * n_docs))
    for i, w in enumerate(VOCAB):
        print(f"  {w:<11}" + "".join(f"{weights[d][i]:>8.3f}" for d in names))
    print("  " + "-" * (11 + 8 * n_docs))
    print(f"  {'norm':<11}" + "".join(
        f"{math.sqrt(sum(x*x for x in weights[d])):>8.4f}" for d in names))
    print("\n  The whole first row is zero. Two documents that share only")
    print("  'the' now share nothing at all.")

    print("\nStage 4. Compare, as an angle.\n")
    raw = {d: [counts[d][w] for w in VOCAB] for d in names}
    print(f"  {'pair':<10}{'cos(raw tf)':>13}{'cos(tf-idf)':>14}")
    print("  " + "-" * 37)
    for a, b in itertools.combinations(names, 2):
        print(f"  {a + '-' + b:<10}{cosine(raw[a], raw[b]):>13.4f}"
              f"{cosine(weights[a], weights[b]):>14.4f}")

    print("\n  Read d2-d3. Raw counts call them 0.57 similar. They share")
    print("  exactly one word, and that word is 'the'. Under tf-idf their")
    print("  cosine is 0.0000, which is the correct answer.")
    print("  Read d1-d2 against d1-d3. Raw counts separate them by a factor")
    print("  of 1.2. tf-idf separates them by a factor of 3.5. The ranking")
    print("  was already right; the weighting made it decisive.")

    print("\nStage 5. Rank documents for a query.")
    print("         score(q, d) = sum over t in q of tf-idf(t, d).\n")
    for query in ("the moon", "moon orbits"):
        qt = query.split()
        scored = sorted(
            ((sum(counts[d][w] * idf[w] for w in qt),
              sum(counts[d][w] for w in qt), d) for d in names),
            key=lambda r: -r[0])
        print(f"  query \"{query}\"")
        print(f"    {'doc':<5}{'raw count':>11}{'tf-idf score':>14}")
        for score, rawc, d in scored:
            print(f"    {d:<5}{rawc:>11}{score:>14.4f}")
    print("\n  On \"the moon\", raw counts tie d1 and d3 at 3 with d2 close")
    print("  behind at 2. tf-idf keeps the tie between d1 and d3, which is")
    print("  right, and sends d2 and d4 to exactly zero, which is also")
    print("  right. Neither of them mentions the moon.")

    print("\n  One honest caveat. In this collection 'a' has df = 1, so it")
    print("  scores the highest idf there is. idf is a statistic of the")
    print("  collection, not a judgement about language. With N = 4 it is")
    print("  measuring almost nothing. Give it 100,000 documents and 'a'")
    print("  falls to zero alongside 'the'.")


# ------------------------------------------------------------------ PMI

def pmi(joint_count, count_w, count_c, n_tokens):
    """log2 of (what we saw) over (what independence predicts).

    0 means exactly independent. Positive means the pair sticks together
    more than chance. Negative means less, and in a sparse matrix negative
    values are mostly noise, which is why PPMI clips them to 0.
    """
    return math.log2((joint_count * n_tokens) / (count_w * count_c))


def ppmi(*args):
    return max(pmi(*args), 0.0)


def demo_pmi():
    N = 1_000_000
    print(f"\nPMI      corpus of N = {N:,} tokens\n")
    rows = [
        # w,      c,        joint, count_w,  count_c
        ("ice",   "cream",    500,   2_000,    2_000),   # strong
        ("ice",   "the",      800,   2_000,   50_000),   # frequent, not strong
        ("data",  "science",   10,   1_000,    1_000),   # modest
        ("cold",  "cream",      4,   2_000,    2_000),   # exactly independent
        ("cream", "asphalt",    1,   2_000,    2_000),   # below chance
    ]
    print(f"  {'pair':<18}{'joint':>7}{'expected':>10}{'PMI':>8}{'PPMI':>7}")
    print("  " + "-" * 50)
    for w, c, joint, cw, cc in rows:
        expected = cw * cc / N          # what independence predicts
        print(f"  {w + ' + ' + c:<18}{joint:>7}{expected:>10.1f}"
              f"{pmi(joint, cw, cc, N):>8.2f}{ppmi(joint, cw, cc, N):>7.2f}")

    print("\n  Row 2 is the one to study. 'ice' sits next to 'the' 800 times,")
    print("  far more often than it sits next to 'cream'. Raw counts would")
    print("  call that the stronger association.")
    print("  PMI does not. 'the' is common on its own, so the expected count")
    print("  is large too, and the ratio collapses from 800 down to 3 bits.")
    print("\n  Row 4 is exactly independent. Joint equals expected, the ratio")
    print("  is 1, and log2(1) = 0. That is the reading to memorise: PMI of")
    print("  zero means the pair tells you nothing.")
    print("\n  Row 5 is below chance, so PMI goes negative. PPMI clips it to 0.")
    print("  Seeing a pair once is no evidence that it is avoided; it is much")
    print("  more likely that the corpus is simply too small to say.")


def main():
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("which", nargs="?", default="all",
                    choices=["tfidf", "corpus", "pmi", "all"])
    a = ap.parse_args()
    if a.which in ("tfidf", "all"):
        demo_tfidf()
    if a.which in ("corpus", "all"):
        demo_corpus()
    if a.which in ("pmi", "all"):
        demo_pmi()
    print()


if __name__ == "__main__":
    main()