"""Prints the six steps of the lesson with real numbers. Run: python3 trace.py"""
import numpy as np

from refusal_toy import D, angle, direction_at, edit_and_measure, forward, measure, orthogonalize, remove, study

LAYER = 5          # read the direction after this layer (0 is before the first layer)
# Layer 5 is a late layer picked by hand; layers 4 to 6 tie on the held-out counts, and 6 reads nearer (6.5 against 8.4 degrees).
SECOND = 0.3       # share of flagged requests that use a second route, for the last line only

model, sets, base = study()
print(f"step 1  read the vectors after layer {LAYER}: {len(sets['train_flagged'])} flagged and {len(sets['train_plain'])} plain requests, {D} numbers each")
u = direction_at(base, LAYER)
print(f"step 2  difference of the two means, length 1: {angle(u, model['routes'][0]):.1f} degrees from the planted direction")
flagged_along = base["test_flagged"][1][LAYER] @ u
plain_along = base["test_plain"][1][LAYER] @ u
cut = (flagged_along.mean() + plain_along.mean()) / 2
sorted_right = int((flagged_along > cut).sum() + (plain_along <= cut).sum())
print(f"step 3  held-out distance along it: flagged {flagged_along.mean():.1f}, plain {plain_along.mean():.1f}; one cut at {cut:.1f} sorts {sorted_right} of 200 right")
hooked = {name: forward(model, sets[name], remove(u)) for name in ("test_flagged", "test_plain")}
length = lambda points: np.linalg.norm(points[-1], axis=1).mean()
print(f"step 4  taken out at all {len(hooked['test_plain'][1])} points: last-layer length of flagged vectors {length(base['test_flagged'][1]):.1f} -> {length(hooked['test_flagged'][1]):.1f}, plain {length(base['test_plain'][1]):.1f} -> {length(hooked['test_plain'][1]):.1f}")
edited_model = orthogonalize(model, u)
weights = {name: forward(edited_model, sets[name])[0] for name in ("test_flagged", "test_plain")}
print(f"step 5  same edit in the weights: largest difference in the answer scores {max(abs(weights[n] - hooked[n][0]).max() for n in weights):.10f}")
result = measure(base, weights)
print(f"step 6  flagged requests refused: {int((base['test_flagged'][0].argmax(1) == 0).sum())} of 100 before, {result['flagged_refused']} of 100 after; "
      f"plain answers changed: {result['changed']} of 100 (mean KL {result['kl']:.3f})")
model2, sets2, base2 = study(second=SECOND)
result2 = edit_and_measure(model2, sets2, base2, direction_at(base2, LAYER))
print(f"limit   with {round(SECOND * 100)} of 100 flagged requests on a second route, the same edit leaves {result2['flagged_refused']} of 100 refused")