"""Print every number the lesson quotes. Run it from this folder: python3 show_numbers.py"""
from prompts import PROMPTS, true_sentence
from toy_lm import (IDK, answer, lines_of, load_corpus, model_with_extra, next_word_probs, ranked, smallest_k,
                    tokenize, train, with_temperature)

corpus = load_corpus()
lines = lines_of(corpus)
model = train(lines)
peru = PROMPTS[0]
PROMPT, TRUTH, SENTENCE = peru["prompt"], peru["truth"], true_sentence(peru)


def table(probs, n=5):
    return ", ".join(f"{w} {p:.1%}" for w, p in ranked(probs)[:n])


print("== The corpus (corpus.txt)")
capitals = [l for l in lines if l.startswith("The capital of")]
print(f"  {len(lines)} sentences, {sum(len(tokenize(l)) for l in lines)} tokens, {len(model['vocab'])} different words")
print(f"  {len(capitals)} are 'The capital of X is Y.' sentences about {len(set(capitals))} countries")
for country in ("France", "Japan", "Italy"):
    print(f"    The capital of {country} appears {sum(1 for l in lines if country in l and l.startswith('The capital'))} times")
print(f"  words after 'is': {sum(model['singles'][('is',)].values())} in all, {len(model['singles'][('is',)])} different")

print(f"\n== Prompt: {PROMPT!r}   (true word: {TRUTH})")
for item in PROMPTS:
    print(f"  {item['name']}: true word {item['truth']!r} in the corpus: {item['truth'] in model['vocab']}; "
          f"{item['prompt'].split()[-2]!r} in the corpus: {item['prompt'].split()[-2] in model['vocab']}")
probs = next_word_probs(model, PROMPT)
print(f"  probabilities add up to {sum(probs.values()):.12f}")
print(f"  top five: {table(probs)}")
print(f"  {TRUTH}: {probs.get(TRUTH, 0):.1%}")
seen = sum(model["singles"][("is",)].values())
kinds = len(model["singles"][("is",)])
trust = seen / (seen + kinds)
print(f"  by hand for Paris: trust {seen}/({seen}+{kinds}) = {trust:.2f}; "
      f"{trust:.2f} * 5/{seen} + {1 - trust:.2f} / {len(model['vocab'])} = {trust * 5 / seen + (1 - trust) / len(model['vocab']):.4f}")
print(f"  ('Peru', 'is') seen in the corpus: {('Peru', 'is') in model['pairs']}")
print(f"  flat guess 1/{len(model['vocab'])} = {1 / len(model['vocab']):.4f}")

one = model_with_extra(lines, SENTENCE, 1)
seen_is, kinds_is = sum(one["singles"][("is",)].values()), len(one["singles"][("is",)])
trust_is, trust_pair = seen_is / (seen_is + kinds_is), 1 / (1 + 1)
lower = trust_is * 1 / seen_is + (1 - trust_is) / len(one["vocab"])
print(f"  by hand for {TRUTH} after k = 1: 'is' now has {seen_is} followers, {kinds_is} different, trust {seen_is}/({seen_is}+{kinds_is}) = {trust_is:.4f}; "
      f"one-word estimate {trust_is:.4f} * 1/{seen_is} + {1 - trust_is:.4f} / {len(one['vocab'])} = {lower:.4f}; "
      f"('Peru', 'is') has 1 follower, 1 different, trust 1/(1+1) = {trust_pair}; "
      f"{trust_pair} * 1 + {1 - trust_pair} * {lower:.4f} = {trust_pair + (1 - trust_pair) * lower:.4f}")

print(f"\n== Adding k copies of {SENTENCE!r}")
for k in range(0, 9):
    probs = next_word_probs(model_with_extra(lines, SENTENCE, k), PROMPT)
    word, p = ranked(probs)[0]
    print(f"  k={k}: top {word} {p:.1%};  {TRUTH} {probs.get(TRUTH, 0):.1%};  runner-up {ranked(probs)[1][0]} {ranked(probs)[1][1]:.1%}")

print("\n== Smallest k")
for thr in (0.0, 0.5, 0.8, 0.9):
    print(f"  top word is {TRUTH} and has at least {thr:.0%}: k = {smallest_k(lines, PROMPT, SENTENCE, TRUTH, thr)}")
print(f"  other questions, no threshold: " + ", ".join(
    f"{i['name']} k = {smallest_k(lines, i['prompt'], true_sentence(i), i['truth'])}" for i in PROMPTS))

print("\n== A fact the corpus states 5 times")
probs = next_word_probs(model, "The capital of France is")
print(f"  top: {table(probs, 3)}")

print("\n== 'I don't know' below a threshold, k = 0")
probs = next_word_probs(model, PROMPT)
for thr in (0.05, 0.10, 0.11, 0.30, 0.50):
    print(f"  threshold {thr:.0%}: {answer(probs, thr)}")
print(f"  threshold 50%, k = 1: {answer(next_word_probs(model_with_extra(lines, SENTENCE, 1), PROMPT), 0.5)}")
print(f"  threshold 50%, France: {answer(next_word_probs(model, 'The capital of France is'), 0.5)}")

print("\n== Temperature, k = 0")
base = next_word_probs(model, PROMPT)
for t in (0.1, 0.2, 0.3, 0.5, 1.0, 1.5, 2.0):
    scaled = with_temperature(base, t)
    word, p = ranked(scaled)[0]
    print(f"  T={t}: top {word} {p:.1%}; {TRUTH} {scaled.get(TRUTH, 0):.1%}; at a 50% threshold the toy says: {answer(scaled, 0.5)}")