"""Choosing K from the singular value spectrum, as Chapter 7 does it.

The singular values tell you how many dimensions your data actually has. Each
one carries a share of the total, and the running total says how much you keep
if you cut at K.

    contribution   r_i = sigma_i / sum(sigma)
    cumulative     R_K = sum(sigma_1..K) / sum(sigma)

    python3 svd_rank.py                 # the book's 8x8 example
    python3 svd_rank.py --squared       # the variance convention
    python3 svd_rank.py --threshold 0.9 # cut wherever you like

The demonstration matrix has two hidden groups, animals and vehicles, sharing
no context. Watch the spectrum discover that without being told.

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"]

# Rows 1-2 and 5-6 are animals, rows 3-4 and 7-8 are vehicles. The two blocks
# share no nonzero column, so the true structure is two dimensional.
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 spectrum(matrix, squared=False):
    """Return (singular values, share of each, cumulative share)."""
    sigma = np.linalg.svd(matrix, compute_uv=False)
    weight = sigma ** 2 if squared else sigma
    share = weight / weight.sum()
    return sigma, share, np.cumsum(share)


def choose_k(cumulative, threshold):
    """Smallest K whose cumulative share reaches the threshold."""
    return int(np.searchsorted(cumulative, threshold) + 1)


def elbow(share):
    """Index after the largest drop between consecutive shares.

    A blunt instrument, and it agrees with the eye on clean spectra. On real
    corpora the curve is smooth and you should use a threshold instead.
    """
    drops = share[:-1] - share[1:]
    return int(np.argmax(drops) + 1)


def main():
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--squared", action="store_true",
                    help="use sigma^2, the explained-variance convention")
    ap.add_argument("--threshold", type=float, default=0.8)
    a = ap.parse_args()

    sigma, share, cum = spectrum(M, a.squared)
    convention = "sigma^2 / sum(sigma^2)" if a.squared else "sigma / sum(sigma)"
    print(f"\n{M.shape[0]}x{M.shape[1]} co-occurrence matrix")
    print(f"contribution measured as {convention}\n")

    k_elbow = elbow(share)
    k_thresh = choose_k(cum, a.threshold)

    print(f"  {'i':>2}{'sigma':>10}{'share':>9}{'cumulative':>12}")
    print("  " + "-" * 33)
    for i, (s, r, c) in enumerate(zip(sigma, share, cum), 1):
        mark = ""
        if i == k_elbow:
            mark += "  <-- elbow"
        if i == k_thresh:
            mark += f"  <-- reaches {a.threshold:.0%}"
        print(f"  {i:>2}{s:>10.3f}{r:>8.1%}{c:>11.1%}{mark}")

    print(f"\n  sum of sigma = {sigma.sum():.3f}")
    print(f"  elbow says          K = {k_elbow}")
    print(f"  {a.threshold:.0%} threshold says  K = {k_thresh}")

    print("\n  This matrix has two hidden groups, animals and vehicles, and")
    print("  they share no context. Nobody told the decomposition that. It")
    print("  read the number of themes off the data.")
    if not a.squared:
        _, sq_share, sq_cum = spectrum(M, squared=True)
        print(f"\n  For comparison, the squared convention gives"
              f" {sq_share[0]:.1%} for the first")
        print(f"  dimension and {sq_cum[1]:.1%} for the first two. Squaring"
              f" exaggerates the")
        print("  lead of the top values. Both conventions are used; say which.")
    print()


if __name__ == "__main__":
    main()