"""Print every number the lesson quotes: the toy's weights for "it", the full grids, the scaling experiment and the cost.

    python3 show_numbers.py
"""
import math

import attention as att

IT = att.TOKENS.index("it")


def show_it():
    print('The word "it" under the first head (what a word points to). d = 2 numbers per query, so scores are divided by sqrt(2) = %.4f' % math.sqrt(2))
    run = att.attend(att.VECTORS, att.HEADS["refers"])
    print("  query of it:", run["queries"][IT])
    print("  %-8s %-10s %-10s %s" % ("word", "key", "score", "weight"))
    for j, t in enumerate(att.TOKENS):
        score = "masked" if j > IT else "%.3f" % run["scores"][IT][j]
        print("  %-8s %-10s %-10s %.4f" % (t, run["keys"][j], score, run["weights"][IT][j]))
    print("  blended note of it:", [round(x, 3) for x in run["outputs"][IT]], "  about:", att.ABOUT)


def show_grids():
    for name, head in att.HEADS.items():
        run = att.attend(att.VECTORS, head)
        print("\nWeights in percent, head %r (%s). Row = the word asking, column = the word it looks at." % (name, att.HEAD_LABELS[name]))
        print("  %-8s" % "" + "".join("%8s" % t for t in att.TOKENS))
        for i, t in enumerate(att.TOKENS):
            print("  %-8s" % t + "".join("%8s" % att.percent(w) if j <= i else "%8s" % "." for j, w in enumerate(run["weights"][i])))
        print("  sentence:", att.view(name, IT)["sentence"], "|", att.view(name, IT)["blend"])
    both = att.multi_head(att.VECTORS, [att.HEADS["refers"], att.HEADS["before"]])
    print("\nTwo heads joined, outputs for it:", [round(x, 3) for x in both[IT]])


def show_scaling():
    print("\nScores of one random query against 8 random keys, 100 trials (numbers of mean 0 and spread 1, seed 6)")
    print("  %-6s %-22s %-22s %s" % ("d", "spread of scores", "biggest weight", "sqrt(d)"))
    print("  %-6s %-10s %-10s %-10s %-10s" % ("", "no divide", "divide", "no divide", "divide"))
    for r in att.scaling_experiment():
        print("  %-6d %-10.2f %-10.2f %-10.2f %-10.2f %.1f" % (r["d"], r["spread_raw"], r["spread_scaled"], r["top_raw"], r["top_scaled"], math.sqrt(r["d"])))


def show_cost():
    print("\nScores one head makes over a prompt, and for %d layers x %d heads" % (att.LAYERS, att.HEADS_PER_LAYER))
    for n in (10, 1000, 2000, 8000):
        print("  %6d words: %14s per head, %18s in all" % (n, f"{att.score_count(n):,}", f"{att.score_count(n) * att.LAYERS * att.HEADS_PER_LAYER:,}"))


if __name__ == "__main__":
    show_it()
    show_grids()
    show_scaling()
    show_cost()