"""A toy network with a refusal planted along one direction, and the steps that find and remove it.

Everything is small and made up here: a residual vector of 32 numbers, six layers, eight possible
answers (answer 0 is "refuse"). No real model, no downloaded weights, only NumPy and a fixed seed.
"""
import numpy as np

D, HIDDEN, LAYERS, ANSWERS = 32, 24, 6, 8     # answer 0 is "refuse"
ROUTES = 2                                    # directions a refusal can travel along
TAU, GAIN, READ = 0.5, 0.5, 2.0               # a layer's refusal unit switches on above TAU and adds GAIN times the excess
FLAG, QUIET = 1.5, 0.1                        # push a flagged request carries; how little plain requests use the routes
SEED, DATA_SEED = 7, 11


def build(seed=SEED):
    """The network: an embedding E, six blocks (W_in, bias, W_out) and a readout U."""
    rng = np.random.default_rng(seed)
    E, _ = np.linalg.qr(rng.normal(size=(D, D)))
    routes = E[:, :ROUTES].T                                  # the planted directions, one unit vector each
    blocks = []
    for _ in range(LAYERS):
        W_in = np.vstack([routes, rng.normal(size=(HIDDEN - ROUTES, D)) / np.sqrt(D)])
        W_out = np.hstack([GAIN * routes.T, 0.15 * rng.normal(size=(D, HIDDEN - ROUTES)) / np.sqrt(HIDDEN)])
        bias = np.r_[np.full(ROUTES, TAU), np.zeros(HIDDEN - ROUTES)]
        blocks.append((W_in, bias, W_out))
    U = np.hstack([READ * routes.sum(0)[:, None], 2 / np.sqrt(D) * rng.normal(size=(D, ANSWERS - 1))])
    return {"E": E, "routes": routes, "blocks": blocks, "U": U}


def requests(n, flagged, rng, noise=1.0, second=0.0):
    """n requests as feature vectors. A flagged one carries a push along route 0, or along route 1 for the first `second` share."""
    f = rng.normal(size=(n, D)) * noise
    f[:, :ROUTES] *= QUIET
    if flagged:
        f[np.arange(n), (np.arange(n) < round(second * n)).astype(int)] += FLAG
    return f


def forward(model, f, hook=lambda x: x):
    """Run requests f. Returns the answer scores and the residual vectors at LAYERS + 1 points."""
    x = hook(f @ model["E"].T)
    points = [x]
    for W_in, bias, W_out in model["blocks"]:
        x = hook(x + np.maximum(x @ W_in.T - bias, 0) @ W_out.T)
        points.append(x)
    return x @ model["U"], points


def refusal_direction(flagged_points, plain_points):
    """Step 2: the difference of the two means, scaled to length 1."""
    d = flagged_points.mean(0) - plain_points.mean(0)
    return d / np.linalg.norm(d)


def remove(u, share=1.0):
    """Step 4: a hook that takes `share` of each vector's part along u out, at every point it is applied."""
    return lambda x: x - share * np.outer(x @ u, u)


def orthogonalize(model, u):
    """Step 5: the same edit written into the weights that write into the residual vector."""
    cut = lambda W: W - np.outer(u, u @ W)
    return {**model, "E": cut(model["E"]), "blocks": [(W_in, bias, cut(W_out)) for W_in, bias, W_out in model["blocks"]]}


def softmax(z):
    e = np.exp(z - z.max(1, keepdims=True))
    return e / e.sum(1, keepdims=True)


def study(noise=1.0, second=0.0, n_train=64, n_test=100):
    """Build the toy, draw the four request sets and run each once unedited."""
    model = build()
    rng = np.random.default_rng(DATA_SEED)
    sets = {name: requests(n, flagged, rng, noise, second) for name, n, flagged in
            (("train_flagged", n_train, True), ("train_plain", n_train, False),
             ("test_flagged", n_test, True), ("test_plain", n_test, False))}
    return model, sets, {name: forward(model, f) for name, f in sets.items()}


def direction_at(base, layer):
    return refusal_direction(base["train_flagged"][1][layer], base["train_plain"][1][layer])


def measure(base, edited_scores):
    """Step 6: held-out flagged requests still refused, plain requests refused, plain answers changed, mean KL on plain."""
    refused = lambda z: int((z.argmax(1) == 0).sum())
    before, after = base["test_plain"][0], edited_scores["test_plain"]
    p, q = softmax(before), softmax(after)
    return {"flagged_refused": refused(edited_scores["test_flagged"]), "plain_refused": refused(after),
            "changed": int((before.argmax(1) != after.argmax(1)).sum()), "kl": float((p * np.log(p / q)).sum(1).mean())}


def edit_and_measure(model, sets, base, u, share=1.0):
    """Remove `share` of direction u at every layer, run the held-out sets, measure."""
    return measure(base, {name: forward(model, sets[name], remove(u, share))[0] for name in ("test_flagged", "test_plain")})


def angle(u, v):
    return float(np.degrees(np.arccos(np.clip(abs(u @ v), 0, 1))))