"""Prints the numbers the lesson text quotes beyond the trace. Run: python3 show_numbers.py"""
import numpy as np

from refusal_toy import study
from show import SHARES, run

default, spread, noisy, quiet = run(), run(second=0.3), run(noise=2.0), run(noise=0.5)

print("where the direction is read (default setting, all of it taken out)")
for layer in range(len(default["angle"])):
    print(f"  after layer {layer}: {default['angle'][layer]:5.1f} degrees off, refused {default['refused'][layer][-1]:3d} of 100, plain answers changed {default['changed'][layer][-1]:2d}, mean KL {default['kl'][layer][-1]:.3f}")
model, _, base = study()
for layer in (0, 5):
    gap = base["train_flagged"][1][layer].mean(0) - base["train_plain"][1][layer].mean(0)
    along = gap @ model["routes"][0]
    print(f"  after layer {layer}: the gap between the two averages is {along:.2f} along the planted direction and {np.linalg.norm(gap - along * model['routes'][0]):.2f} across it")
print("how much is taken out (read after layer 5)")
for i, share in enumerate(SHARES):
    print(f"  {share:3d}%: refused {default['refused'][5][i]:3d} of 100, plain answers changed {default['changed'][5][i]:2d}")
for name, r in (("a second route for 30 of 100", spread), ("high noise", noisy), ("low noise", quiet)):
    print(f"{name} (read after layer 5, all taken out): {r['angle'][5]:.1f} degrees off, refused {r['refused'][5][-1]} of 100, plain answers changed {r['changed'][5][-1]}, mean KL {r['kl'][5][-1]:.3f}")