"""Print every number the lesson quotes. Run: python3 show_numbers.py"""
import speculative_sim as sim

big = sim.big


def main() -> None:
    plain, draft = sim.plain_seconds(), sim.draft_pass().seconds
    print(f"big model {big.PARAMS:,} weights, {big.FP16_BYTES * big.PARAMS / 1e9:.2f} GB; draft model {sim.DRAFT_PARAMS:,} weights, {big.FP16_BYTES * sim.DRAFT_PARAMS / 1e9:.2f} GB")
    print(f"plain big pass {plain * 1000:.3f} ms; draft pass {draft * 1000:.3f} ms = {100 * draft / plain:.1f}% of a big pass (c = {draft / plain:.4f}); ridge {big.ridge():.2f}")
    print()
    print("a big pass over n inputs (the guesses plus one):")
    for n in (1, 3, 5, 9, 100, 153, 154, 200, 1000):
        p = big.prefill(n)
        print(f"  n={n:5,}: {p.seconds * 1000:8.3f} ms, {p.bound}")
    print()
    print("one round at k = 2, 4, 8:", {k: f"{sim.round_seconds(k) * 1000:.3f} ms" for k in (2, 4, 8)})
    print()
    print("tokens per big pass and speed-up at k = 4 (guesses kept with chance a, independently):")
    for a in (0.3, 0.5, 0.6, 0.8, 0.9, 0.95):
        print(f"  a={a}: {sim.expected_tokens(a, 4):.2f} tokens per pass, {sim.speedup(a, 4):.2f} times as fast, simulated {sim.simulate_tokens_per_pass(a, 4, 100_000, seed=4):.2f}")
    print("break-even acceptance:", {k: round(sim.break_even(k), 3) for k in (1, 2, 4, 8)})
    print("best k and its speed-up:", {a: (sim.best_k(a), round(sim.speedup(a, sim.best_k(a)), 2)) for a in (0.5, 0.6, 0.8, 0.9, 0.95)})
    print("speed-up at a = 0.6, 0.8, 0.9 for k = 1..12:")
    for a in (0.6, 0.8, 0.9):
        print(f"  a={a}:", [round(sim.speedup(a, k), 2) for k in range(1, 13)])
    print()
    print("many users sharing the chip (the check covers users * (k + 1) inputs):")
    print("  users when the check stops being free:", {k: sim.users_when_checking_stops_being_free(k) for k in (2, 4, 6, 8)})
    print("  speed-up at a = 0.8, k = 4:", {u: round(sim.speedup(0.8, 4, users=u), 2) for u in (1, 16, 30, 31, 32, 64, 128, 154, 1000)})
    print("  users when it stops paying (a = 0.8):", {k: sim.users_when_it_stops_paying(0.8, k) for k in (2, 4, 8)})
    print()
    print(f"toy pair: draft agrees with the big model at {float(sim.agreement()) * 100:.1f}% of places ({sim.agreement()}); measured against the formula at that chance")
    for k in (1, 2, 4, 8):
        print(f"  k={k}: measured {sim.toy_tokens_per_pass(k):.2f} tokens per pass, formula {sim.expected_tokens(float(sim.agreement()), k):.2f}")
    print(f"  reply: {' '.join(sim.REPLY)}")
    print()
    p, q = sim.HAND_P, sim.HAND_Q
    print("hand example (big model p, small model q):", {x: f"{p[x]} / {q[x]}" for x in p})
    print("  chance a guess x is drawn and kept:", {x: str(q[x] * sim.keep_chance(p[x], q[x])) for x in p}, "total kept", sim.keep_total(p, q))
    print("  redraw after a rejection:", {x: str(v) for x, v in sim.residual(p, q).items()}, "-> output", {x: str(v) for x, v in sim.output_distribution(p, q).items()})
    print()
    print("widget sentences:")
    for pct, k in ((sim.OPENING["percent"], sim.OPENING["k"]), (20, 4), (50, 4), (100, 8)):
        v = sim.view(pct, k, sim.DRAFT_STOPS_MS[sim.DEFAULT_DRAFT_STOP], sim.OPENING["seed"])
        print(f"  {pct}% k={k}: {v['spec']} {v['verdict']} | {v['average']}")
    print(f"  plain: {sim.view(80, 4, sim.DRAFT_STOPS_MS[2], sim.OPENING['seed'])['plain']}")
    print(f"  draft stops (ms): {[round(x, 3) for x in sim.DRAFT_STOPS_MS]}; opening run rounds {sim.plan_run(sim.OPENING['percent'], sim.OPENING['k'], sim.OPENING['seed'])}")


if __name__ == "__main__":
    main()