"""COALS, traced as Chapter 7 traces it.

Counts become correlations, negatives are discarded, positives are square
rooted. Run this and you get the same three tables the lectures show for the
woodchuck corpus.

    python3 coals.py              # all three steps
    python3 coals.py --step 2     # just the correlation matrix
    python3 coals.py --cell a if  # the arithmetic for one cell, shown long

The point of the correlation step is that it factors out frequency. Watch the
word "a", which is the most common word here, stop dominating.

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

import argparse
import math

WORDS = ["a", "as", "chuck", "could", "how", "if", "much", "wood",
         "woodchuck", "would", ",", ".", "?"]

# The step-1 count matrix from the lectures, built with a ramped window of 4
# over: "How much wood would a woodchuck chuck, if a woodchuck could chuck
# wood? As much wood as a woodchuck would, if a woodchuck could chuck wood."
COUNTS = [
    [0,  5,  9,  6,  1, 10,  4,  8, 18,  9, 10,  0,  0],   # a
    [5,  4,  2,  1,  0,  0,  7, 10,  3,  2,  1,  0,  5],   # as
    [9,  2,  0,  8,  0,  0,  1,  9, 11,  2,  4,  3,  3],   # chuck
    [6,  1,  8,  0,  0,  4,  0,  6,  8,  0,  2,  2,  2],   # could
    [1,  0,  0,  0,  0,  4,  3,  0,  0,  2,  0,  0,  0],   # how
    [10, 0,  5,  4,  0,  0,  4,  3, 10,  3,  0,  0,  0],   # if
    [4,  7,  1,  0,  4,  0,  0, 10,  2,  3,  0,  0,  3],   # much
    [8, 10,  9,  6,  3,  0, 10,  2,  8,  5,  0,  4,  6],   # wood
    [18, 3, 11,  0,  0, 10,  2,  8,  8, 10,  1,  1,  1],   # woodchuck
    [9,  2,  2,  0,  2,  3,  5,  8,  0,  5,  0,  0,  0],   # would
    [10, 1,  4,  2,  0,  0,  0, 10,  5,  0,  0,  0,  0],   # ,
    [0,  0,  3,  2,  0,  0,  4,  1,  1,  0,  0,  0,  0],   # .
    [0,  5,  3,  2,  0,  0,  3,  6,  1,  0,  0,  0,  0],   # ?
]


def margins(M):
    """Row sums, column sums, and the grand total T."""
    rows = [sum(r) for r in M]
    cols = [sum(M[i][j] for i in range(len(M))) for j in range(len(M[0]))]
    return rows, cols, sum(rows)


def correlation(w_ab, row_a, col_b, T):
    """Equation 7.7.

    Numerator: what we observed (scaled by T) minus what independence predicts
    (the product of the margins). Same comparison PMI makes.

    Denominator: how much each margin could vary. This is what puts every cell
    on the same [-1, 1] scale regardless of how common the words are, and it
    is exactly what HAL lacked.
    """
    num = T * w_ab - row_a * col_b
    den = math.sqrt(row_a * (T - row_a) * col_b * (T - col_b))
    return num / den if den else 0.0


def coals_value(r):
    """Equation 7.8. Discard negatives, damp the survivors."""
    return math.sqrt(r) if r > 0 else 0.0


def table(M, title, fmt="{:>7.3f}"):
    print(f"\n{title}\n")
    head = "  " + " " * 10 + "".join(f"{w[:6]:>7}" for w in WORDS)
    print(head)
    for i, w in enumerate(WORDS):
        print(f"  {w:<10}" + "".join(fmt.format(v) for v in M[i]))


def one_cell(a, b):
    """Show the arithmetic for a single cell, the long way."""
    rows, cols, T = margins(COUNTS)
    i, j = WORDS.index(a), WORDS.index(b)
    w, ra, cb = COUNTS[i][j], rows[i], cols[j]
    num = T * w - ra * cb
    den = math.sqrt(ra * (T - ra) * cb * (T - cb))
    r = num / den
    print(f"\nCELL ({a}, {b})\n")
    print(f"  count w[{a}][{b}]        = {w}")
    print(f"  row sum for '{a}'         = {ra}")
    print(f"  column sum for '{b}'      = {cb}")
    print(f"  grand total T            = {T}")
    print(f"\n  numerator   = T*w - row*col = {T}*{w} - {ra}*{cb}"
          f" = {T*w} - {ra*cb} = {num}")
    print(f"  denominator = sqrt({ra}*{T-ra} * {cb}*{T-cb}) = {den:.0f}")
    print(f"\n  r = {num} / {den:.0f} = {r:+.3f}")
    print(f"  after clipping and square root: {coals_value(r):.3f}")
    if w == 0:
        print(f"\n  Note the count is zero and r is negative. '{a}' and '{b}'")
        print("  do not merely fail to co-occur, they avoid each other.")


def main():
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--step", type=int, choices=[1, 2, 3], default=None)
    ap.add_argument("--cell", nargs=2, metavar=("A", "B"))
    a = ap.parse_args()

    if a.cell:
        one_cell(*a.cell)
        print()
        return

    rows, cols, T = margins(COUNTS)
    R = [[correlation(COUNTS[i][j], rows[i], cols[j], T)
          for j in range(len(WORDS))] for i in range(len(WORDS))]
    C = [[coals_value(v) for v in row] for row in R]

    if a.step in (None, 1):
        table(COUNTS, "STEP 1: raw counts, ramped window of 4", "{:>7d}")
        print(f"\n  grand total T = {T}")
        print(f"  row sums    {dict(zip(WORDS, rows))}")
    if a.step in (None, 2):
        table(R, "STEP 2: counts converted to correlations")
        print("\n  Every value now sits in [-1, 1], whatever the word's"
              " frequency.")
        print("  'a' is the most frequent word here and no longer dominates.")
    if a.step in (None, 3):
        table(C, "STEP 3: negatives set to 0, positives square rooted")
        neg = sum(1 for row in R for v in row if v < 0)
        print(f"\n  {neg} of {len(WORDS)**2} cells were negative and are now 0,"
              f" which keeps the matrix sparse.")
        print("  The square root pulls in the large values, so a few strong")
        print("  pairings cannot dominate a vector.")
    print()


if __name__ == "__main__":
    main()