"""Contextual embeddings, worked as Chapter 15 works them.

Three things, all reproducing the chapter's tables:

    the averaging cost   what one vector per word does to 'bank'
    ELMo layer mixing    the softmax weights, and what each layer contributes
    masked prediction    why BERT masks 15 per cent and not 100

    python3 contextual.py           # all three
    python3 contextual.py --bank    # static against contextual
    python3 contextual.py --elmo    # layer weights, three task profiles
    python3 contextual.py --mask    # the masking budget

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

import argparse
import math

# A toy space on three axes: (money, river, generic). The two senses of 'bank'
# are given separately, as a contextual model would produce them.
SPACE = {
    "bank/money":  (0.90, 0.05, 0.20),
    "bank/river":  (0.05, 0.90, 0.20),
    "deposit":     (0.95, 0.02, 0.15),
    "loan":        (0.92, 0.03, 0.18),
    "vault":       (0.88, 0.01, 0.10),
    "water":       (0.02, 0.93, 0.15),
    "erosion":     (0.01, 0.90, 0.08),
    "flood":       (0.03, 0.95, 0.12),
}


def cosine(u, v):
    d = sum(a * b for a, b in zip(u, v))
    n = math.sqrt(sum(a * a for a in u)) * math.sqrt(sum(b * b for b in v))
    return d / n if n else 0.0


def neighbours(vec, exclude=(), top=4):
    scored = [(w, cosine(vec, v)) for w, v in SPACE.items() if w not in exclude]
    return sorted(scored, key=lambda p: -p[1])[:top]


# ------------------------------------------------------------- the averaging

def show_bank():
    money = SPACE["bank/money"]
    river = SPACE["bank/river"]
    static = tuple((a + b) / 2 for a, b in zip(money, river))

    print("\nSTEP 1  what one vector per word costs\n")
    print("  Two senses of 'bank', as a contextual model would give them:\n")
    print(f"    bank (money)  {money}")
    print(f"    bank (river)  {river}")
    print(f"\n  A static embedding has to pick one vector. Trained on a corpus")
    print(f"  with both senses in equal measure, it lands on their average:\n")
    print(f"    bank (static) ({static[0]:.3f}, {static[1]:.3f},"
          f" {static[2]:.3f})")

    print(f"\n  Now ask each of the three for its nearest neighbours.\n")
    for label, vec, ex in [("bank, money sense", money, ("bank/money",)),
                           ("bank, river sense", river, ("bank/river",)),
                           ("bank, static average", static, ())]:
        top = neighbours(vec, exclude=ex + ("bank/money", "bank/river"))
        cells = "   ".join(f"{w} {s:+.3f}" for w, s in top)
        print(f"  {label:<22}{cells}")

    print(f"\n  The two sense vectors give clean, single-topic lists. The")
    print(f"  average gives a list that mixes both and commits to neither.")

    print(f"\n  Measure it. Cosine of the static average against each sense:\n")
    print(f"    to the money sense  {cosine(static, money):.4f}")
    print(f"    to the river sense  {cosine(static, river):.4f}")
    print(f"    the two senses to each other  {cosine(money, river):.4f}")

    print(f"\n  The average sits {cosine(static, money):.2f} from each sense,")
    print(f"  which sounds close until you notice the senses are only")
    print(f"  {cosine(money, river):.2f} from each other. The static vector is")
    print(f"  equally wrong in two directions rather than right in one.")

    print(f"\n  And the damage scales with how balanced the corpus is. A")
    print(f"  corpus that is 90 per cent money sense gives:\n")
    for share in (0.5, 0.7, 0.9, 0.99):
        v = tuple(share * a + (1 - share) * b for a, b in zip(money, river))
        print(f"    {share:>5.0%} money   to money {cosine(v, money):.4f}"
              f"   to river {cosine(v, river):.4f}")
    print(f"\n  So a static vector does not represent a word. It represents")
    print(f"  the word's frequency-weighted mixture in whatever corpus you")
    print(f"  happened to train on.")


# --------------------------------------------------------------- ELMo mixing

def show_elmo():
    print("\n\nSTEP 2  ELMo mixes the layers, and the task picks the mix\n")
    print("  A biLM of L layers gives L+1 representations per token. ELMo")
    print("  combines them with softmax-normalised task weights:\n")
    print("    ELMo_k = gamma * sum_j  s_j * R_kj\n")
    print("  Only s_0..s_L and gamma are learned downstream. The biLM is")
    print("  frozen, which is why ELMo was cheap to adopt.\n")

    layers = ["layer 0, characters", "layer 1, lower biLSTM",
              "layer 2, upper biLSTM"]
    profiles = {
        "part-of-speech tagging": [0.2, 1.8, 0.5],
        "word sense disambiguation": [0.1, 0.6, 2.0],
        "no preference": [1.0, 1.0, 1.0],
    }

    for task, raw in profiles.items():
        m = max(raw)
        e = [math.exp(r - m) for r in raw]
        s = [x / sum(e) for x in e]
        print(f"  {task}")
        print(f"    {'layer':<24}{'raw weight':>12}{'s_j softmax':>14}")
        print("    " + "-" * 50)
        for name, r, w in zip(layers, raw, s):
            print(f"    {name:<24}{r:>12.2f}{w:>14.4f}")
        print(f"    {'':<24}{'':>12}{sum(s):>14.4f}\n")

    print("  The three profiles are the finding of the ELMo paper, stated as")
    print("  weights. Lower layers carry more syntax, upper layers more")
    print("  meaning, and which one a task wants is a fact you can read off")
    print("  the learned s_j rather than a thing you decide in advance.")
    print("\n  That is also the strongest evidence that depth is doing")
    print("  something structured rather than just adding capacity.")


# ------------------------------------------------------------- masking budget

def show_mask():
    print("\n\nSTEP 3  why mask 15 per cent\n")
    print("  BERT trains by hiding tokens and predicting them. The masking")
    print("  rate is a trade, and both ends of it are easy to see.\n")

    n = 512
    print(f"  A sequence of {n} tokens.\n")
    print(f"  {'rate':>7}{'masked':>9}{'visible':>9}"
          f"{'training signal':>18}{'context left':>15}")
    print("  " + "-" * 60)
    for rate in (0.01, 0.15, 0.50, 0.90):
        masked = int(n * rate)
        print(f"  {rate:>6.0%}{masked:>9}{n - masked:>9}"
              f"{masked:>18}{1 - rate:>14.0%}")

    print("\n  At 1 per cent you get 5 predictions from a forward pass that")
    print("  cost you the whole sequence. Training is correct and slow.")
    print("\n  At 90 per cent you get plenty of predictions and almost no")
    print("  context to make them from. The task becomes guessing.")
    print("\n  BERT chose 15 per cent, which gives 76 gradients per sequence")
    print("  with 85 per cent of the sentence still readable. Later work has")
    print("  found higher rates workable on larger models, so the number is")
    print("  a balance point rather than a constant of nature.")

    print("\n\n  THE PRETRAIN AND FINE-TUNE MISMATCH\n")
    print("  The token [MASK] appears in every pretraining batch and never")
    print("  once at fine-tuning time. A model that keyed on it would break")
    print("  the moment it was used for real.")
    print("\n  So of the 15 per cent chosen, BERT does this:\n")
    chosen = 100
    split = [("replaced with [MASK]", 0.80),
             ("replaced with a random token", 0.10),
             ("left unchanged", 0.10)]
    print(f"  {'treatment':<32}{'share':>8}{'per 100 chosen':>17}")
    print("  " + "-" * 58)
    for name, share in split:
        print(f"  {name:<32}{share:>7.0%}{share * chosen:>17.0f}")

    print("\n  The last two rows are the interesting ones. Because a chosen")
    print("  position might be unchanged, the model cannot tell which")
    print("  positions were chosen. So it has to build a usable")
    print("  representation of every token, not just the masked ones.")
    print("\n  The random-token row does something else. It forces the model")
    print("  to notice when a token does not fit its context, which is a")
    print("  skill no amount of [MASK] prediction would teach.")


def main():
    ap = argparse.ArgumentParser(description=__doc__)
    for f in ("bank", "elmo", "mask"):
        ap.add_argument(f"--{f}", action="store_true")
    a = ap.parse_args()
    picked = a.bank or a.elmo or a.mask
    if a.bank or not picked:
        show_bank()
    if a.elmo or not picked:
        show_elmo()
    if a.mask or not picked:
        show_mask()
    print()


if __name__ == "__main__":
    main()