"""Print every number the lesson quotes. Run it from this folder: python3 show_numbers.py"""
from collections import Counter

import sampler as s

toy = s.toy
lines = toy.lines_of(toy.load_corpus())
model = toy.train(lines)
tea, plant = "How do I make tea?", "How do I water a plant?"
wait = s.as_chances(*s.row_of(model, ["Wait"]))
tokens = [t for t, _ in wait]


def pct(chances):
    return ", ".join(f"{t} {100 * c:.1f}%" for t, c in chances)


print("== The Wait row (what followed the token Wait in the 17 lines)")
ranked, total = s.row_of(model, ["Wait"])
print(f"  {total} times: " + ", ".join(f"{t} {c}" for t, c in ranked) + f"; {sum('Wait' in toy.tokenize(l) for l in lines)} of {len(lines)} lines have it")
for t in (0.1, 0.2, 0.5, 1.0, 2.0):
    print(f"  temperature {t}: {pct(s.with_temperature(wait, t))}")

print("\n== Top-p on the Wait row at temperature 1 (kept tokens, rescaled)")
for p in (0.5, 0.7, 0.8, 0.9, 1.0):
    print(f"  p {p}: {pct(s.reshape(wait, 1.0, p))}")
print("  temperature first, then top-p:")
for t, p in ((0.2, 0.7), (0.1, 0.7), (2.0, 0.7)):
    print(f"  temperature {t}, p {p}: {pct(s.reshape(wait, t, p))}")
print(f"  top-p first, then temperature 0.2 (p 0.7): {pct(s.with_temperature(s.top_p(wait, 0.7), 0.2))}")

print(f"\n== {s.DRAWS} draws from seed 1 (counts for {', '.join(tokens)})")
for t, p in ((1.0, 1.0), (0.2, 1.0), (0.1, 1.0), (2.0, 1.0), (1.0, 0.7)):
    picks, _ = s.draw(s.reshape(wait, t, p), s.start_state(1))
    print(f"  temperature {t}, p {p}: {s.tally(picks, tokens)}")
picks, state = s.draw(wait, s.start_state(1))
again, _ = s.draw(wait, state)
print(f"  temperature 1, p 1, seed 1, first press: {' '.join(picks)}")
print(f"  ... and the next press: {s.tally(again, tokens)}")
print(f"  same seed, again: {s.tally(s.draw(wait, s.start_state(1))[0], tokens)}; seed 2: {s.tally(s.draw(wait, s.start_state(2))[0], tokens)}")

print("\n== Worked example, pass 11 of tea (cup 3, mug 2)")
row = s.as_chances(*s.row_of(model, ["bag", "in", "a"]))
for t in (1.0, 0.5, 2.0):
    powered = [(tok, c ** (1 / t)) for tok, c in row]
    print(f"  temperature {t}: power {1 / t:g}: " + ", ".join(f"{tok} {c:.4f}" for tok, c in powered)
          + f"; total {sum(c for _, c in powered):.4f}; chances {pct(s.with_temperature(row, t))}")
for u in (0.5, 0.65):
    print(f"  u {u}: " + "; ".join(f"temperature {t} takes {s.choose(s.with_temperature(row, t), u)}" for t in (1.0, 0.5, 2.0)))
print("  tea, pass 1 (Boil 4, Heat 1): " + "; ".join(f"temperature {t}: {pct(s.with_temperature(s.as_chances(*s.row_of(model, ['make', 'tea', '?'])), t))}" for t in (0.5, 1.0, 2.0)))

print("\n== The generator")
raw = toy.seed_step(toy.seed_state(1))[1]
print(f"  first number for seed 1 with no warm-up: {raw:.4e}; with {s.WARMUP} thrown away: {toy.seed_step(s.start_state(1))[1]:.4f}")
heat = lambda start: sum(s.generate(model, tea, start(seed), cap=1)[0] == ["Heat"] for seed in range(1000))
print(f"  Heat as the first token over seeds 0-999: {heat(toy.seed_state)} of 1000 with no warm-up, {heat(s.start_state)} with it (its chance is 20%)")

print("\n== 1,000 replies from one seed (seed 1): different replies, of those not among the toy's own lines, top reply count")
for name, prompt in (("tea", tea), ("plant", plant)):
    known = s.corpus_replies(lines, prompt)
    for t, p in ((1.0, 1.0), (0.5, 1.0), (0.2, 1.0), (0.05, 1.0), (2.0, 1.0), (1.0, 0.7)):
        counted = Counter(s.reply_text(r) for r in s.replies(model, prompt, 1000, 1, t, p))
        print(f"  {name} temperature {t}, p {p}: {len(counted)} different, {sum(1 for x in counted if x not in known)} new, top reply {counted.most_common(1)[0][1]}")
    print(f"  {name}: {len(known)} different replies among its {sum(l.startswith(prompt) for l in lines)} lines")
print("  greedy tea: " + toy.reply_of(toy.generate(model, tea, 40)))
first = Counter(r[0] for r in s.replies(model, tea, 1000, 1))
print(f"  tea first token over those 1,000 replies: {dict(first)}")

print("\n== Two runs part ways (plant, seed 1, temperature 1)")
runs = s.replies(model, plant, 6, 1)
dry, wet = runs[0], runs[5]
fork = s.first_split(dry, wet)
print(f"  they share the first {fork} tokens, then {dry[fork]!r} ({len(dry) - fork - 1} tokens from there on, end mark not counted) and {wet[fork]!r} ({len(wet) - fork - 1})")
print(f"  next-token row at the fork: {s.row_of(model, ['the', 'soil', 'is'])}")
print("  one: " + s.reply_text(dry))
print("  two: " + s.reply_text(wet))
shared = set(s.reply_text(dry[fork:]).split(". ")) & set(s.reply_text(wet[fork:]).split(". "))
print(f"  sentences the two have in common after the fork: {sorted(shared)}")

print("\n== Frequency check: 100,000 draws from seed 1")
for t, p in ((1.0, 1.0), (0.5, 1.0), (2.0, 0.8)):
    chances = s.reshape(wait, t, p)
    picks, _ = s.draw(chances, s.start_state(1), 100000)
    seen = Counter(picks)
    worst = max(abs(seen[tok] / 100000 - c) for tok, c in chances)
    print(f"  temperature {t}, p {p}: largest gap between share drawn and chance {100 * worst:.2f} points; kept tokens {len(chances)}")