"""Guards for refusal_toy.py: the direction found against the planted one, the hook against the weight edit, the story the lesson tells, the limit case, and that every
number the lesson, its trace and its figures print is the one the code gives."""
import io
import re
import subprocess
import sys
import xml.etree.ElementTree as ET
from contextlib import redirect_stdout
from functools import lru_cache
from pathlib import Path

import numpy as np

import make_figs
import show
import refusal_toy as toy
from refusal_toy import angle, direction_at, edit_and_measure, forward, orthogonalize, remove, study

HERE = Path(__file__).resolve().parent
LESSON = HERE.parent
SOURCE = (HERE / "refusal_toy.py").read_text(encoding="utf-8")
LAYER = 5


@lru_cache(maxsize=None)
def toy_study(second=0.0):
    model, sets, base = study(second=second)
    return model, sets, base, direction_at(base, LAYER)


def refused_mask(model, f, hook=lambda x: x):
    return forward(model, f, hook)[0].argmax(1) == 0


def test_direction_found_is_near_the_planted_one_and_sorts_new_requests():
    """Guards the method finding something other than the planted direction: the difference of means must land near it, and one cut along it must sort held-out requests."""
    model, sets, base, u = toy_study()
    assert angle(u, model["routes"][0]) < 10 and abs(np.linalg.norm(u) - 1) < 1e-12
    flagged, plain = base["test_flagged"][1][LAYER] @ u, base["test_plain"][1][LAYER] @ u
    cut = (flagged.mean() + plain.mean()) / 2
    assert int((flagged > cut).sum() + (plain <= cut).sum()) == 200


def test_hook_and_weight_edit_agree_at_every_point_and_for_every_share_cut():
    """Guards the claim that cutting u from the weights is the same edit as taking it out of the vector: the vectors and the answer scores agree at all 7 points."""
    model, sets, base, u = toy_study()
    edited = orthogonalize(model, u)
    for name in ("test_flagged", "test_plain"):
        hooked_scores, hooked_points = forward(model, sets[name], remove(u))
        scores, points = forward(edited, sets[name])
        assert len(points) == toy.LAYERS + 1 and abs(scores - hooked_scores).max() < 1e-9
        assert all(abs(a - b).max() < 1e-9 and abs(b @ u).max() < 1e-9 for a, b in zip(points, hooked_points))
    W = model["blocks"][LAYER][2]
    assert abs(orthogonalize(model, u)["blocks"][LAYER][2].T @ u).max() < 1e-12 and abs(W.T @ u).max() > 0.4      # the cut removes a part that was there


def test_the_story_reproduces():
    """Guards the lesson's claims: every flagged request refused before and none after, one plain answer in 100 changed, refusals falling in a narrow band of the share
    taken out, a later layer finding the direction better, and a second route that one direction does not remove."""
    model, sets, base, u = toy_study()
    assert int(refused_mask(model, sets["test_flagged"]).sum()) == 100 and not refused_mask(model, sets["test_plain"]).any()
    full = edit_and_measure(model, sets, base, u)
    assert full["flagged_refused"] == 0 and full["changed"] == 1 and full["plain_refused"] == 0 and 0.005 < full["kl"] < 0.007
    by_share = {share: edit_and_measure(model, sets, base, u, share / 100)["flagged_refused"] for share in (0, 10, 15, 20, 25, 30)}
    assert by_share == {0: 100, 10: 100, 15: 90, 20: 44, 25: 5, 30: 0}
    at0 = edit_and_measure(model, sets, base, direction_at(base, 0))
    assert at0["flagged_refused"] == 13 and at0["changed"] == 8 and 37.9 == round(angle(direction_at(base, 0), model["routes"][0]), 1)
    angles = [angle(direction_at(base, k), model["routes"][0]) for k in range(toy.LAYERS + 1)]
    assert angles == sorted(angles, reverse=True)
    model2, sets2, base2, u2 = toy_study(0.3)
    left = refused_mask(model2, sets2["test_flagged"], remove(u2))
    assert int(left.sum()) == 29 and int(left[:30].sum()) == 29 and not left[30:].any()      # all that stay refused are on the second route
    assert int(refused_mask(model2, sets2["test_flagged"]).sum()) == 100


def test_route_part_follows_the_layer_rule_in_the_lesson():
    """Guards the sentence in Go deeper that a refusal unit multiplies the part along the route by 1.5 and takes 0.25 off: the toy's averages must stay near it."""
    model, sets, base, u = toy_study()
    x, rule = toy.FLAG, []
    for _ in range(toy.LAYERS + 1):
        rule.append(x)
        x = (1 + toy.GAIN) * x - toy.GAIN * toy.TAU
    measured = [float((p @ model["routes"][0]).mean()) for p in base["test_flagged"][1]]
    assert all(abs(a - b) < 0.2 for a, b in zip(rule, measured)) and [round(v, 2) for v in rule[:3]] == [1.5, 2.0, 2.75]
    assert [round(v, 1) for v in rule[3:]] == [3.9, 5.6, 8.1, 11.9]


def reading():
    prose = (LESSON / "lesson.md").read_text(encoding="utf-8")
    return prose, {name: (LESSON / "figs" / name).read_text(encoding="utf-8") for name in make_figs.FIGURES}


def gaps(base, model, layer):
    """The gap between the two training averages along the planted direction and across it."""
    d = base["train_flagged"][1][layer].mean(0) - base["train_plain"][1][layer].mean(0)
    along = float(d @ model["routes"][0])
    return along, float(np.sqrt(d @ d - along ** 2))


def test_every_number_the_lesson_and_its_figures_quote_comes_from_the_code():
    """Guards prose, trace and figure numbers drifting from the code: each fragment is built from a value computed here and must appear in the lesson or the figure."""
    prose, figures = reading()
    model, sets, base, u = toy_study()
    grid = show.grid()[1]                                    # normal noise: one run for each second-route share
    run, second30 = grid[0], grid[2]
    i = lambda share: show.SHARES.index(share)
    ang = lambda k: run["angle"][k]
    full = edit_and_measure(model, sets, base, u)
    flagged, plain = base["test_flagged"][1][LAYER] @ u, base["test_plain"][1][LAYER] @ u
    cut = (flagged.mean() + plain.mean()) / 2
    length = lambda name: np.linalg.norm(base[name][1][-1], axis=1).mean()
    hooked = {n: forward(model, sets[n], remove(u)) for n in ("test_flagged", "test_plain")}
    hooked_length = lambda name: np.linalg.norm(hooked[name][1][-1], axis=1).mean()
    a0, c0 = gaps(base, model, 0)
    a5, c5 = gaps(base, model, LAYER)
    lines = len(SOURCE.splitlines())
    L5, L0 = run["refused"][LAYER][i(100)], run["refused"][0][i(100)]
    wanted = {"lesson.md": [
        "toy network of 100 lines of Python", f"`refusal_toy.py` is {lines} lines", f"All {run['base']} flagged requests",
        f"and {full['changed']} of 100 plain answers changes", f"0 of 100 flagged requests are refused, down from {run['base']}", "none is, while 99 of 100 plain answers stay the same",
        f"Share taken out to 20%: {run['refused'][LAYER][i(20)]} flagged requests are still refused", "between 10% and 30%",
        f"the direction is found {ang(0)} degrees off the planted one, {L0} flagged requests are still refused and {run['changed'][0][i(100)]} plain answers change",
        f"leaves {second30['refused'][LAYER][i(100)]} refused", f"Run {len(sets['train_flagged'])} flagged and {len(sets['train_plain'])} plain requests through the network",
        f"keep each one's {toy.D} numbers after layer {LAYER}", f"It lands {ang(LAYER)} degrees from the planted one", f"{len(sets['test_flagged'])} flagged requests the step never saw sit at {flagged.mean():.1f} on average",
        f"100 plain ones at {plain.mean():.1f}", f"One cut at {cut:.1f} sorts all {int((flagged > cut).sum() + (plain <= cut).sum())}", f"{len(hooked['test_plain'][1])} times",
        f"Flagged requests refused: {run['base']} before, {full['flagged_refused']} after. Plain requests whose answer changes: {full['changed']} of 100",
        f"The direction is found from {len(sets['train_flagged'])} and {len(sets['train_plain'])} requests and lands {ang(LAYER)} degrees", f"Flagged requests sit at {flagged.mean():.1f} along the direction and plain ones at {plain.mean():.1f}",
        f"falls from {length('test_flagged'):.1f} to {hooked_length('test_flagged'):.1f} and plain ones stay at {hooked_length('test_plain'):.1f}", f"with a mean KL of {full['kl']:.3f}",
        "nothing happens at 10%, and by 30% no flagged request is refused",
        f"Put 30 of the 100 flagged requests on a second route, and the same edit leaves {second30['refused'][LAYER][i(100)]} refused",
        f"{ang(0)} degrees off, {L0} flagged requests are still refused, and {run['changed'][0][i(100)]} plain answers change",
        f"gave {ang(LAYER)} degrees, {L5} and {run['changed'][LAYER][i(100)]}", "while all 100 refusals went",
        f"{toy.HIDDEN} hidden units", f"`TAU = {toy.TAU}`", f"`GAIN = {toy.GAIN}` times", f"`READ = {toy.READ:.0f}` times", f"`FLAG = {toy.FLAG}`",
        "1.5, 2.0, 2.75, 3.9, 5.6, 8.1 and 11.9", f"After layer {LAYER} the gap is {a5:.2f} along r and {c5:.2f} across it, which is {ang(LAYER)} degrees",
        f"After layer 0 it is {a0:.2f} along and {c0:.2f} across, which is {ang(0)} degrees", f"the 7 points",
        f"gap along it is {a5:.2f} against {c5:.2f} across it",
        f"differ by {a0:.2f} along the planted direction but by {c0:.2f} across it", f"their scores differ by {0:.10f}",
        f"{full['flagged_refused']} of 100 flagged requests are refused and {full['changed']} of 100 plain answers changes",
        f"Taking out 10% of the direction leaves {run['refused'][LAYER][i(10)]} of 100 flagged requests refused, and 30% leaves none",
        f"{ang(0)} degrees off the planted one. After layer 5 it is {ang(LAYER)}",
        f"the edit leaves {second30['refused'][LAYER][i(100)]} refused",
        f"The direction found then sits {second30['angle'][LAYER]:.1f} degrees off the planted one", f"layer 6 sits nearer the planted direction, at {ang(6)} degrees",
    ]}
    wanted.update({
        "the-map.svg": [f"{len(sets['test_flagged'][:30])} flagged requests", f"a cut at {cut:.1f}", "along the direction"],
        "weight-edit.svg": [f"{toy.HIDDEN} bars", f"all {toy.HIDDEN} columns: 0.00", f"the other {toy.HIDDEN - 1}"],
        "share-removed.svg": [f"{run['refused'][LAYER][i(15)]} at 15%", f"{run['refused'][LAYER][i(20)]} at 20%", f"{run['refused'][LAYER][i(25)]} at 25%", f"{run['refused'][LAYER][i(10)]} at 10%", "0 from 30% on"],
        "where-to-read.svg": [f"layer {k}: {ang(k)} degrees off, {run['refused'][k][i(100)]} refused, {run['changed'][k][i(100)]} changed" for k in range(toy.LAYERS + 1)],
        "second-route.svg": [f"{second30['refused'][LAYER][i(100)]} are refused, {second30['refused'][LAYER][i(100)]} of them among the 30", f">{second30['refused'][LAYER][i(100)]} refused<"],
    })
    texts = {"lesson.md": prose, **figures}
    missing = [f"{name}: {fragment!r}" for name, fragments in wanted.items() for fragment in fragments if fragment not in texts[name]]
    assert not missing, "not found: " + "; ".join(missing)
    assert lines < 150 and run["refused"][LAYER][i(30)] == 0


def test_the_figures_are_the_ones_the_code_draws_and_follow_the_house_style():
    """Guards a figure going stale or off-palette: the files on disk equal what make_figs draws, the lesson shows each one, and each parses, has a title and a description,
    a white ground, no script, no web address, only the house colours and no text under 20 units."""
    prose, figures = reading()
    assert set(figures) == set(re.findall(r"\]\(figs/([\w-]+\.svg)\)", prose)) == set(make_figs.FIGURES)
    allowed = {make_figs.INK, make_figs.GREY, make_figs.ORANGE, make_figs.GREEN, make_figs.FILL, make_figs.PANEL, "#ffffff"}
    web_address = re.compile("ht" + r"tps?://")
    for name, svg in figures.items():
        assert svg == make_figs.FIGURES[name](), name
        root = ET.fromstring(svg)
        assert root.find("{%s}title" % make_figs.NAMESPACE) is not None and root.find("{%s}desc" % make_figs.NAMESPACE) is not None and 'fill="#ffffff"' in svg, name
        assert "<script" not in svg and not web_address.search(svg.replace(make_figs.NAMESPACE, "")) and root.get("width") == "600", name
        assert {c.lower() for c in re.findall(r"#[0-9a-fA-F]{6}\b", svg)} <= {c.lower() for c in allowed}, name
        assert all(int(size) >= 20 for size in re.findall(r'font-size="(\d+)"', svg)), name
    for path in HERE.glob("*.py"):
        assert not web_address.search(path.read_text(encoding="utf-8").replace("ht" + "tp://www.w3.org", "")), path.name


def table_rows(prose):
    return re.findall(r"^\| \d\. [^|]+\| `(\w+\.py):(\d+)`<br>`(.+?)` \|", prose, re.M)


def test_code_blocks_are_slices_of_the_real_file_and_the_table_points_at_its_lines():
    """Guards a `file:line` in the code table pointing at the wrong line, and a code block in Go deeper that no longer matches the file."""
    prose, _ = reading()
    blocks = re.findall(r"```python\n(.*?)\n```", prose, re.S)
    assert len(blocks) == 4 and all(block in SOURCE for block in blocks)
    rows = table_rows(prose)
    assert [(name, int(n)) for name, n, _ in rows] == [("refusal_toy.py", 45), ("refusal_toy.py", 51), ("trace.py", 14), ("refusal_toy.py", 57), ("refusal_toy.py", 62), ("refusal_toy.py", 87)]
    for name, number, code in rows:
        line = (HERE / name).read_text(encoding="utf-8").splitlines()[int(number) - 1]
        assert re.sub(r"\s+# .*$", "", line).strip() == code, (name, number)


def run_trace(source):
    out = io.StringIO()
    with redirect_stdout(out):
        exec(compile(source, "trace.py", "exec"), {"__name__": "__main__"})
    return out.getvalue()


def test_trace_is_the_block_in_the_lesson_and_the_your_turn_answer_is_right():
    """Guards the pasted run drifting from `python3 trace.py` (the real command, from the code folder) and the Your turn answer: changing line 6 to LAYER = 0 must give
    the lines in the details block."""
    prose, _ = reading()
    real = subprocess.run([sys.executable, "-B", "trace.py"], cwd=HERE, capture_output=True, text=True, check=True).stdout
    pasted = re.search(r"python3 trace\.py\n```\n\nPhone readers: slide the box sideways to see every column\.\n\n```text\n(.*?)\n```", prose, re.S).group(1)
    assert real == pasted + "\n"
    trace = (HERE / "trace.py").read_text(encoding="utf-8")
    assert trace.splitlines()[5] == "LAYER = 5          # read the direction after this layer (0 is before the first layer)" and "change line 6 to `LAYER = 0`" in prose
    changed = run_trace(trace.replace("LAYER = 5 ", "LAYER = 0 ", 1))
    answer = re.search(r"<summary>What you should see</summary>.*?```text\n(.*?)\n```", prose, re.S).group(1)
    assert answer + "\n" == changed and "13 of 100 after" in answer
    assert [a == b for a, b in zip(real.splitlines(), changed.splitlines())] == [False, False, False, False, True, False, False]