#!/usr/bin/env python
"""Lab 0: weight-space programming on a toy layer small enough to print.

    uv run --with numpy --with matplotlib python labs/lab0_toy.py      # seconds, any laptop

An MLP's last matrix, down_proj, is W (here 4 x 6; in Qwen3.5-0.8B it is 1,024 x 3,584). Neuron j fires with
strength a[j] and writes its column W[:, j] into the stream, so the layer adds y = W a = sum_j a[j] * W[:, j].
The paper's edit picks a direction u and some neurons S and sets  W' = W - alpha * u (u^T W[:, S])  on those
columns. Every claim below is checked with an assert, so if this script finishes, they all hold.
"""
import os

import numpy as np

np.set_printoptions(precision=3, suppress=True)
rng = np.random.default_rng(0)
d, n = 4, 6
W = rng.normal(size=(d, n))
u = rng.normal(size=d)
u /= np.linalg.norm(u)            # the direction we care about (in the paper: "toward looping")
S = [1, 4]                        # the neurons we allow ourselves to touch (in the paper: 64 of 3,584)
others = [j for j in range(n) if j not in S]


def edit(W, u, alpha, S=None):
    g = u @ W                     # g[j] = how much neuron j writes along u, per unit of activation
    if S is not None:
        mask = np.zeros(W.shape[1])
        mask[S] = 1
        g = g * mask
    return W - alpha * np.outer(u, g)


print("1. How much each neuron writes along u (u . W[:, j]):", u @ W)

print("\n2. The edit scales the chosen neurons' push along u by (1 - alpha) and leaves everything else alone")
P = np.eye(d) - np.outer(u, u)    # removes the u part of any vector
for alpha, name in [(0, "no edit"), (1, "erase"), (2, "mirror"), (6, "flip and amplify 5x (the paper)"),
                    (-4, "amplify 5x the other way")]:
    W2 = edit(W, u, alpha, S)
    before, after = (u @ W)[S], (u @ W2)[S]
    assert np.allclose(after, (1 - alpha) * before)
    assert np.allclose(W2[:, others], W[:, others])     # untouched neurons stay identical
    assert np.allclose(P @ W2, P @ W)                     # nothing changes outside the direction u
    print(f"   alpha={alpha:>2} ({name:31s}) push along u of neurons {S}: {before} -> {after}")

print("\n3. The change is rank one: one direction out (u), one pattern in (u^T W on the chosen columns)")
D = edit(W, u, 6, S) - W
sv = np.linalg.svd(D, compute_uv=False)
assert sv[1] < 1e-12 * sv[0]
print("   singular values of W' - W:", sv)

print("\n4. The weight edit is steering: W' a = W a - alpha * s * u, with s = u . (W[:, S] a[S])")
alpha = 6
for name, a in [("chosen neurons quiet", np.where(np.isin(np.arange(n), S), 0.0, rng.normal(size=n))),
                ("chosen neurons firing", np.where(np.isin(np.arange(n), S), 3.0, rng.normal(size=n)))]:
    s = u @ (W[:, S] @ a[S])
    y_edit, y_steer = edit(W, u, alpha, S) @ a, W @ a - alpha * s * u
    assert np.allclose(y_edit, y_steer)
    print(f"   {name:22s} s = {s:+.3f}   change to the output |W'a - Wa| = {np.linalg.norm(y_edit - W @ a):.3f}")
print("   Same weights, different tokens, different push: the edit only acts when the chosen neurons fire.")

import matplotlib  # noqa: E402
matplotlib.use("Agg")
import matplotlib.pyplot as plt  # noqa: E402

fig, ax = plt.subplots(1, 2, figsize=(11, 3.6))
im = ax[0].imshow(edit(W, u, 6, S) - W, cmap="RdBu", vmin=-6, vmax=6)
ax[0].set_title("W' - W: only the chosen columns change")
ax[0].set_xlabel("neuron (column)")
ax[0].set_ylabel("stream dimension (row)")
fig.colorbar(im, ax=ax[0])
al = np.linspace(-5, 8, 100)
ax[1].plot(al, 1 - al, color="#0969da")
for x, lab in [(0, "no edit"), (1, "erase"), (2, "mirror"), (6, "the paper")]:
    ax[1].plot(x, 1 - x, "o", color="#cf222e")
    ax[1].annotate(lab, (x, 1 - x), xytext=(6, 4), textcoords="offset points")
ax[1].axhline(0, color="black", lw=0.5)
ax[1].set_xlabel("alpha")
ax[1].set_ylabel("push along u after / before")
ax[1].set_title("What alpha does to the chosen neurons")
fig.tight_layout()
out = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "results", "figures")
os.makedirs(out, exist_ok=True)
fig.savefig(os.path.join(out, "lab0_toy.png"), dpi=130)
print("\nall checks passed; figure in results/figures/lab0_toy.png")