"""The SVD itself: factorise, truncate, and see what the truncation cost.

`svd_rank.py` answers "how big should K be". This one answers "what actually
happens when I cut at K". They are the two halves of Chapter 7's second repair.

    python3 svd.py            # factorise, truncate at every K, show the error
    python3 svd.py --k 2      # dense word vectors at K=2, and similarities
    python3 svd.py --ppmi     # run PPMI over the counts first, as LSA does

Watch for the moment the truncation stops losing information and starts
removing noise.

Install:  pip install numpy
"""

import argparse

try:
    import numpy as np
except ImportError:
    raise SystemExit("this one needs numpy:  pip install numpy")

WORDS = ["dog", "cat", "car", "truck", "pet", "vet", "drive", "fuel"]
CONTEXTS = ["pet", "feed", "fur", "vet", "drive", "fuel", "road", "tyre"]

M = np.array([
    [4, 3, 3, 2, 0, 0, 0, 0],   # dog
    [4, 3, 4, 2, 0, 0, 0, 0],   # cat
    [0, 0, 0, 0, 4, 3, 3, 2],   # car
    [0, 0, 0, 0, 4, 3, 4, 2],   # truck
    [3, 2, 2, 3, 0, 0, 0, 0],   # pet
    [2, 1, 1, 4, 0, 0, 0, 0],   # vet
    [0, 0, 0, 0, 3, 2, 3, 2],   # drive
    [0, 0, 0, 0, 2, 3, 2, 3],   # fuel
], dtype=float)


def ppmi(counts):
    """Positive pointwise mutual information, Chapter 7 Equations 7.2 and 7.3."""
    total = counts.sum()
    p_wc = counts / total
    p_w = p_wc.sum(axis=1, keepdims=True)
    p_c = p_wc.sum(axis=0, keepdims=True)
    with np.errstate(divide="ignore", invalid="ignore"):
        pmi = np.log2(p_wc / (p_w * p_c))
    return np.nan_to_num(np.maximum(pmi, 0.0), neginf=0.0, posinf=0.0)


def truncate(U, S, Vt, k):
    """Rebuild the matrix from only the top k singular values. Equation 7.6."""
    return U[:, :k] @ np.diag(S[:k]) @ Vt[:k, :]


def cosine(u, v):
    n = np.linalg.norm(u) * np.linalg.norm(v)
    return float(u @ v / n) if n else 0.0


def main():
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--k", type=int, default=None,
                    help="show dense word vectors at this K")
    ap.add_argument("--ppmi", action="store_true",
                    help="reweight with PPMI before factorising, as LSA does")
    a = ap.parse_args()

    A = ppmi(M) if a.ppmi else M
    label = "PPMI-weighted counts" if a.ppmi else "raw counts"
    U, S, Vt = np.linalg.svd(A)

    print(f"\nMATRIX  {A.shape[0]} words x {A.shape[1]} contexts, {label}")
    print(f"rank    {np.linalg.matrix_rank(A)}"
          f"   (nonzero singular values: {(S > 1e-10).sum()})")

    print(f"\nFACTORS   A = U . diag(sigma) . V^T")
    print(f"  U  {U.shape}   sigma  {S.shape}   V^T  {Vt.shape}")
    print(f"  sigma = " + "  ".join(f"{s:.3f}" for s in S))

    print(f"\nTRUNCATION  what each K costs\n")
    print(f"  {'K':>2}{'kept share':>12}{'error':>10}{'stored':>9}"
          f"{'vs full':>9}")
    print("  " + "-" * 42)
    full_store = A.size
    for k in range(1, len(S) + 1):
        Ak = truncate(U, S, Vt, k)
        err = np.linalg.norm(A - Ak) / np.linalg.norm(A)
        share = S[:k].sum() / S.sum()
        store = k * (A.shape[0] + A.shape[1] + 1)
        print(f"  {k:>2}{share:>11.1%}{err:>10.3f}{store:>9}"
              f"{store / full_store:>8.1f}x")

    print("\n  'error' is the relative Frobenius distance from the original.")
    print("  It falls fast to K=2, then only crawls. Where it flattens, the")
    print("  extra dimensions were describing noise, not structure.")
    print("\n  'stored' counts the numbers you must keep: k*(rows+cols+1).")
    print("  Note it passes 1.0x at K=4. Truncation only saves space while K")
    print("  stays small, which on a real vocabulary it always does.")

    if a.k:
        k = a.k
        dense = U[:, :k] @ np.diag(S[:k])
        print(f"\nDENSE WORD VECTORS at K={k}"
              f"   ({A.shape[1]} dimensions down to {k})\n")
        for w, v in zip(WORDS, dense):
            print(f"  {w:<8}" + "  ".join(f"{x:+7.3f}" for x in v))

        print(f"\nCOSINE SIMILARITY in the compressed space\n")
        pairs = [("dog", "cat"), ("car", "truck"), ("dog", "car"),
                 ("pet", "vet"), ("drive", "fuel")]
        for x, y in pairs:
            i, j = WORDS.index(x), WORDS.index(y)
            print(f"  {x:<7} {y:<7} {cosine(dense[i], dense[j]):+.3f}")
        print("\n  Words from the same group score 1. Words from different")
        print("  groups score 0. Two dimensions were enough, because the data")
        print("  only ever had two, which is what svd_rank.py detects.")
    print()


if __name__ == "__main__":
    main()