"""word2vec by hand, as Chapter 9 works it: negative sampling and one update.

Three things, all reproducing the book's tables:

    the 0.75 exponent      why negatives are drawn from count^0.75
    one training step      sigmoid, prediction minus target, update
    the analogy            king - man + woman, and the exclusion trap

    python3 word2vec.py                  # all three
    python3 word2vec.py --exponent 1.0   # sample from raw counts instead
    python3 word2vec.py --exponent 0.5   # flatten it further
    python3 word2vec.py --no-exclude     # let query words win the analogy
    python3 word2vec.py --apple          # apple - iphone, and corpus bias
    python3 word2vec.py --subtract mac   # strip a different association

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

import argparse
import math

COUNTS = {"the": 1000, "of": 600, "dog": 50, "cat": 40, "aardvark": 2}

# A second toy space, on axes (technology, fruit, brand), standing in for a
# space trained on technology news. 'apple' carries a big technology component
# and a smaller fruit one, which is a fact about the corpus, not about English.
TECH_SPACE = {
    "apple":   (0.85, 0.45, 0.80),
    "iphone":  (0.95, 0.02, 0.85),
    "mac":     (0.90, 0.03, 0.75),
    "samsung": (0.88, 0.02, 0.80),
    "android": (0.92, 0.01, 0.70),
    "orange":  (0.10, 0.88, 0.25),
    "banana":  (0.02, 0.95, 0.05),
    "fruit":   (0.03, 0.96, 0.02),
    "juice":   (0.06, 0.86, 0.10),
}

# A toy space with interpretable axes: (royal, male, female).
SPACE = {
    "king":   (1.00, 0.90, 0.10),
    "man":    (0.10, 0.90, 0.10),
    "woman":  (0.10, 0.10, 0.90),
    "queen":  (0.95, 0.15, 0.85),
    "prince": (0.90, 0.85, 0.15),
    "throne": (0.80, 0.45, 0.45),
    "child":  (0.20, 0.50, 0.50),
}


def sigmoid(z):
    return 1 / (1 + math.exp(-z))


def dot(a, b):
    return sum(x * y for x, y in zip(a, b))


def cosine(u, v):
    n = math.sqrt(dot(u, u)) * math.sqrt(dot(v, v))
    return dot(u, v) / n if n else 0.0


# ------------------------------------------------- the sampling distribution

def sampling_distribution(counts, exponent):
    """Equation 8.1. Raising counts to a power below 1 lifts the tail."""
    total = sum(counts.values())
    p = {w: c / total for w, c in counts.items()}
    raised = {w: c ** exponent for w, c in counts.items()}
    z = sum(raised.values())
    q = {w: v / z for w, v in raised.items()}
    return p, raised, q


def demo_sampling(exponent):
    p, raised, q = sampling_distribution(COUNTS, exponent)
    print(f"\nNEGATIVE SAMPLING DISTRIBUTION   exponent = {exponent}\n")
    print(f"  {'word':<10}{'count':>7}{'p(w)':>9}{'count^e':>10}"
          f"{'q(w)':>9}{'q/p':>9}")
    print("  " + "-" * 54)
    for w in COUNTS:
        print(f"  {w:<10}{COUNTS[w]:>7}{p[w]:>9.4f}{raised[w]:>10.1f}"
              f"{q[w]:>9.4f}{q[w]/p[w]:>8.2f}x")

    hi, lo = max(COUNTS, key=COUNTS.get), min(COUNTS, key=COUNTS.get)
    print(f"\n  '{hi}' is sampled {q[hi]/p[hi]:.2f}x its share,"
          f" '{lo}' {q[lo]/p[lo]:.2f}x its share.")
    if exponent == 1.0:
        print("  At exponent 1.0 nothing is bent. Nearly every negative would")
        print("  be a stopword, which teaches the model almost nothing.")
    elif exponent < 0.6:
        print("  Flattened this far, rare words appear as negatives far more")
        print("  often than they ever appear as real context words.")
    else:
        print("  This is the setting word2vec ships with. It trims the head")
        print("  and lifts the tail, without going all the way to uniform.")


# ------------------------------------------------------- one training update

def demo_update(eta=0.1):
    v = (0.5, -0.2)            # centre word, 'cat'
    pairs = [("sat", (0.4, 0.1), 1), ("the", (-0.3, 0.6), 0)]

    print(f"\n\nONE TRAINING STEP   centre 'cat' v = {v},  eta = {eta}\n")
    print(f"  {'pair':<22}{'z = v.u':>10}{'sigma(z)':>11}{'label':>7}"
          f"{'gradient':>11}")
    print("  " + "-" * 61)
    grads, updated = {}, {}
    for name, u, label in pairs:
        z = dot(v, u)
        s = sigmoid(z)
        g = s - label                      # prediction minus target
        grads[name] = (u, g, s)
        kind = "positive" if label else "negative"
        print(f"  {'cat, ' + name + ' (' + kind + ')':<22}{z:>+10.3f}"
              f"{s:>11.4f}{label:>7}{g:>+11.4f}")

    print("\n  A negative gradient pulls the pair together.")
    print("  A positive gradient pushes it apart.")

    print(f"\n  {'vector':<8}{'update':<40}{'new value'}")
    print("  " + "-" * 72)
    for name, (u, g, _) in grads.items():
        new = tuple(round(x - eta * g * y, 4) for x, y in zip(u, v))
        updated[name] = new
        arith = f"{u} - {eta}({g:+.4f}){v}"
        print(f"  {name:<8}{arith:<40}{new}")

    # the centre word collects the gradient from every pair
    v_new = tuple(v[i] - eta * sum(g * u[i] for u, g, _ in grads.values())
                  for i in range(len(v)))

    print(f"\n  DID IT WORK?   centre vector also moved to"
          f" {tuple(round(x,4) for x in v_new)}\n")
    print(f"  {'score':<28}{'before':>10}{'after':>10}   want")
    print("  " + "-" * 60)
    for name, (u, _, before) in grads.items():
        after = sigmoid(dot(v_new, updated[name]))
        want = "up" if name == "sat" else "down"
        print(f"  {'sigma(cat . ' + name + ')':<28}{before:>10.4f}"
              f"{after:>10.4f}   {want}")


# ------------------------------------------------------------- the analogy

def demo_analogy(exclude=True):
    a, b, c = "man", "king", "woman"          # analogy(a, b, c) = v_b - v_a + v_c
    target = tuple(SPACE[b][i] - SPACE[a][i] + SPACE[c][i] for i in range(3))

    print(f"\n\nANALOGY   v_{b} - v_{a} + v_{c}"
          f" = ({target[0]:.2f}, {target[1]:.2f}, {target[2]:.2f})\n")
    print(f"  {'word':<9}{'royal':>7}{'male':>7}{'female':>8}{'cosine':>10}")
    print("  " + "-" * 46)
    ranked = sorted(((w, cosine(target, v)) for w, v in SPACE.items()),
                    key=lambda p: -p[1])
    for w, s in ranked:
        v = SPACE[w]
        note = "   <- query word" if w in (a, b, c) else ""
        print(f"  {w:<9}{v[0]:>7.2f}{v[1]:>7.2f}{v[2]:>8.2f}{s:>10.4f}{note}")

    pool = [w for w, _ in ranked if not exclude or w not in (a, b, c)]
    print(f"\n  answer: {pool[0]}"
          f"   ({'query words excluded' if exclude else 'NO EXCLUSION'})")
    if exclude:
        cheat = ranked[0][0]
        if cheat in (a, b, c):
            print(f"  Without the exclusion '{cheat}' would have won, and the")
            print("  analogy would look correct while proving nothing.")
        else:
            order = [w for w, _ in ranked]
            best_q = min((w for w in (a, b, c) if w in order),
                         key=order.index)
            rank = order.index(best_q) + 1
            print(f"  Here the best query word, '{best_q}', only reaches rank"
                  f" {rank}. On real")
            print("  spaces it usually places first, because the target stays")
            print("  close to v_c. Forget the exclusion and your scores inflate.")


def demo_apple(strip="iphone", target="apple"):
    """Subtract one word from another and see which sense survives."""
    def rank(vec, exclude=()):
        return sorted(((w, cosine(vec, v)) for w, v in TECH_SPACE.items()
                       if w not in exclude), key=lambda p: -p[1])

    print(f"\n\nNEIGHBOURS OF '{target}'   (space trained on technology news)\n")
    print(f"  {'neighbour':<10}{'cosine':>9}")
    print("  " + "-" * 21)
    for w, c in rank(TECH_SPACE[target], exclude=(target,))[:6]:
        print(f"  {w:<10}{c:>+9.4f}")
    print(f"\n  The top of that list is technology. '{target}' is a company")
    print("  in this corpus, and the other sense is buried underneath.")

    d = tuple(TECH_SPACE[target][i] - TECH_SPACE[strip][i]
              for i in range(len(TECH_SPACE[target])))
    print(f"\n  vec({target}) - vec({strip}) = "
          f"({d[0]:+.2f}, {d[1]:+.2f}, {d[2]:+.2f})\n")
    print(f"  {'neighbour':<10}{'cosine':>9}")
    print("  " + "-" * 21)
    for w, c in rank(d, exclude=(target, strip))[:6]:
        print(f"  {w:<10}{c:>+9.4f}")

    print(f"\n  Subtracting cancels what the two words share and keeps what")
    print(f"  only '{target}' has. The technology component nearly annihilates.")
    print("  The buried sense comes to the surface, and the words that were")
    print("  nearest before now point the other way.")


def main():
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--exponent", type=float, default=0.75)
    ap.add_argument("--no-exclude", action="store_true",
                    help="let the query words compete in the analogy")
    ap.add_argument("--apple", action="store_true",
                    help="only the apple minus iphone demonstration")
    ap.add_argument("--subtract", default="iphone",
                    help="which word to strip from 'apple'")
    a = ap.parse_args()
    if a.apple or a.subtract != "iphone":
        demo_apple(strip=a.subtract)
        print()
        return
    demo_sampling(a.exponent)
    demo_update()
    demo_analogy(exclude=not a.no_exclude)
    demo_apple()
    print()


if __name__ == "__main__":
    main()