"""Draw the five figures of lesson 32 from refusal_toy.py, so that no number in a figure is typed by hand.

    python3 make_figs.py          rewrite ../figs/*.svg
    python3 make_figs.py --check  exit 1 if a figure on disk differs (test_refusal_toy.py runs the same check)

All five figures are 600 units wide. Text is at least 22 units (20 for axis ticks) in ink or grey on white or a light fill; orange and green are used on marks
only, and every mark that colour tells apart also differs in shape or dash and is named in words.
"""
import sys
from pathlib import Path
from xml.sax.saxutils import escape

import numpy as np

import show
from refusal_toy import LAYERS, direction_at, forward, orthogonalize, remove, study

FIGS = Path(__file__).resolve().parent.parent / "figs"
INK, GREY, ORANGE, GREEN, FILL, PANEL = "#0a1626", "#44566b", "#b36200", "#00855F", "#d2dae4", "#ebf0f6"
DASH = "10 8"
LAYER, SHOWN = 5, 30                    # the layer the lesson reads; requests drawn per set in figure 1


def num(x: float) -> str:
    return f"{x:.2f}".rstrip("0").rstrip(".")


# The SVG namespace is a name that every standalone .svg file must carry, not an address this program visits. It is written in two pieces so that a plain
# search of code/ for the start of a web address finds nothing (test_refusal_toy.py runs that search).
NAMESPACE = "ht" + "tp://www.w3.org/2000/svg"


def frame(width: int, height: int, title: str, desc: str, body: list[str], defs: str = "") -> str:
    head = (f'<svg xmlns="{NAMESPACE}" viewBox="0 0 {width} {height}" width="{width}" height="{height}" role="img" '
            f'font-family="system-ui, sans-serif" font-size="22" fill="{GREY}">')
    return "\n".join([head, f"  <title>{escape(title)}</title>", f"  <desc>{escape(desc)}</desc>", *([f"  {defs}"] if defs else []),
                      '  <rect width="100%" height="100%" fill="#ffffff"/>', *("  " + line for line in body), "</svg>", ""])


def text(x: float, y: float, content: str, anchor: str = "start", color: str = INK, bold: bool = False, size: int = 22) -> str:
    extra = (f' text-anchor="{anchor}"' if anchor != "start" else "") + (' font-weight="bold"' if bold else "") + (f' font-size="{size}"' if size != 22 else "")
    return f'<text x="{num(x)}" y="{num(y)}"{extra} fill="{color}">{escape(content)}</text>'


def triangle(x: float, y: float) -> str:
    return f'<path d="M{num(x)} {num(y - 8)}L{num(x + 8)} {num(y + 6)}L{num(x - 8)} {num(y + 6)}z" fill="{ORANGE}" stroke="#ffffff" stroke-width="1"/>'


def circle(x: float, y: float, r: float = 6, color: str = GREEN) -> str:
    return f'<circle cx="{num(x)}" cy="{num(y)}" r="{num(r)}" fill="{color}" stroke="#ffffff" stroke-width="1"/>'


def square(x: float, y: float, half: float = 6) -> str:
    return f'<rect x="{num(x - half)}" y="{num(y - half)}" width="{num(2 * half)}" height="{num(2 * half)}" fill="#ffffff" stroke="{GREEN}" stroke-width="3"/>'


def tick(x: float, y: float, label: str, anchor: str = "middle") -> str:
    return text(x, y, label, anchor, GREY, size=20)


# --- figure 1: the map, before and after ------------------------------------------------------------------------------------------

def fig_the_map() -> str:
    model, sets, base = study()
    u = direction_at(base, LAYER)
    points = {name: base[name][1][LAYER][:SHOWN] for name in ("test_flagged", "test_plain")}
    both = np.vstack(list(points.values()))
    across = both - np.outer(both @ u, u)                                   # what is left of each vector once its part along u is taken away
    v = np.linalg.svd(across - across.mean(0), full_matrices=False)[2][0]    # the direction of the most spread that remains
    v = v if v[np.abs(v).argmax()] > 0 else -v
    along = {name: p @ u for name, p in points.items()}
    height = {name: (p - np.outer(p @ u, u)) @ v for name, p in points.items()}
    cut = (base["test_flagged"][1][LAYER] @ u).mean() / 2 + (base["test_plain"][1][LAYER] @ u).mean() / 2
    scale, low, mid = 21.5, -1.5, (min(h.min() for h in height.values()) + max(h.max() for h in height.values())) / 2
    top, floor = 76, 76 + 187
    body = []
    for k, (title, left) in enumerate((("Before", 20), ("After", 310))):
        px = lambda x: left + 6 + (x - low) * scale
        py = lambda y: top + 93.5 - (y - mid) * scale
        body.append(text(left, 32, title, bold=True))
        body.append(f'<path d="M{num(px(0))} {top}V{floor}" stroke="{FILL}" stroke-width="2"/>')
        body.append(f'<path d="M{left} {floor}H{left + 270}" stroke="{GREY}" stroke-width="2"/>')
        for t in (0, 4, 8):
            body.append(tick(px(t), floor + 28, str(t)))
        body.append(text(left + 135, floor + 62, "along the direction", "middle", GREY))
        body.append(f'<text x="{left - 2}" y="170" text-anchor="middle" transform="rotate(-90 {left - 2} 170)" fill="{GREY}">other spread</text>')
        if k == 0:
            body.append(f'<path d="M{num(px(cut))} {top}V{floor}" stroke="{GREY}" stroke-width="3" stroke-dasharray="{DASH}"/>')
            body.append(text(px(cut), 62, "cut", "middle"))
        for x, y in zip(along["test_plain"], height["test_plain"]):
            body.append(circle(px(x if k == 0 else 0.0), py(y)))
        for x, y in zip(along["test_flagged"], height["test_flagged"]):
            body.append(triangle(px(x if k == 0 else 0.0), py(y)))
    body.append(text(20 + 6 + (0 - low) * scale, 62, "plain", "middle"))
    body.append(text(20 + 6 + (float(np.mean(along["test_flagged"])) - low) * scale, 62, "flagged", "middle"))
    body.append(text(316, 62, "both sets"))
    flagged, plain = along["test_flagged"], along["test_plain"]
    desc = (f"Two panels of {SHOWN} flagged requests (orange triangles) and {SHOWN} plain requests (green circles) at layer {LAYER}. Across: how far along the found direction a "
            f"request sits. Up, labelled other spread: the biggest spread that remains. Before: flagged requests sit between {flagged.min():.1f} and {flagged.max():.1f}, plain requests between "
            f"{plain.min():.1f} and {plain.max():.1f}, and a cut at {cut:.1f} sorts them. After the direction is taken out, every request sits at 0.0 along it, "
            f"both sets in one stripe, and their up positions are unchanged.")
    return frame(600, floor + 80, "Before the cut flagged and plain requests sit apart along one direction; after it they sit together", desc, body)


# --- figure 2: the same edit in the weights --------------------------------------------------------------------------------------

def fig_weight_edit() -> str:
    model, sets, base = study()
    u = direction_at(base, LAYER)
    W = model["blocks"][LAYER][2]                                           # the matrix that writes into the vector after the layer read
    before = W.T @ u                                                        # how much of each of its 24 columns lies along u
    after = orthogonalize(model, u)["blocks"][LAYER][2].T @ u
    others = before[1:]
    pitch, scale = 540 / len(before), 200.0
    body = []
    for title, y, values in (("Before: the part of each column along u", 170, before), ("After: the same, once that part is cut", 306, after)):
        body.append(text(30, 34 if y == 170 else 272, title, bold=True))
        body.append(f'<path d="M30 {y}H570" stroke="{GREY}" stroke-width="2"/>')
        for k, value in enumerate(values):
            if abs(value) >= 0.0005:
                body.append(f'<rect x="{num(30 + k * pitch + 3)}" y="{num(y - max(value, 0) * scale)}" width="{num(pitch - 6)}" height="{num(abs(value) * scale)}" '
                            f'fill="{ORANGE if k == 0 else INK}"/>')
    body.append(text(30 + pitch + 10, 100, f"refusal unit: {before[0]:.2f}"))
    body.append(text(30, 212, f"the other {len(others)}: {others.min():.2f} to {others.max():.2f}", color=GREY))
    body.append(text(30, 344, f"all {len(after)} columns: {np.abs(after).max():.2f}", color=GREY))
    body.append(text(30, 386, "bars: one column of the matrix each", color=GREY))
    desc = (f"Two rows of {len(before)} bars, one for each column of the matrix that writes into the vector after layer {LAYER}, giving how much of the column lies along u. "
            f"Before, column 0, the refusal unit, has {before[0]:.2f} along u and the other {len(others)} columns between {others.min():.2f} and {others.max():.2f}. "
            f"After the edit every column has {np.abs(after).max():.2f} along u.")
    return frame(600, 406, "Cutting the part along u from each column that writes into the vector changes almost one column", desc, body)


# --- figure 3: how much to take out -----------------------------------------------------------------------------------------------

def fig_share_removed() -> str:
    r = show.run()
    refused, changed = r["refused"][LAYER], r["changed"][LAYER]
    x0, x1, top, base = 84, 560, 60, 320
    px = lambda share: x0 + share / 100 * (x1 - x0)
    py = lambda count: base - count / 100 * (base - top)
    body = [text(20, 36, "requests out of 100", color=GREY)]
    for tick_value in (0, 50, 100):
        body.append(f'<path d="M{x0} {num(py(tick_value))}H{x1}" stroke="{FILL}" stroke-width="2"/>')
        body.append(tick(x0 - 8, py(tick_value) + 7, str(tick_value), "end"))
    body.append(f'<path d="M{x0} {base}H{x1}" stroke="{GREY}" stroke-width="2"/>')
    for share in (0, 25, 50, 75, 100):
        body.append(tick(px(share), base + 30, f"{share}%"))
    body.append(text((x0 + x1) / 2, base + 64, "share of the direction taken out", "middle", GREY))
    body.append(f'<path d="M{" L".join(f"{num(px(s))} {num(py(c))}" for s, c in zip(show.SHARES, changed))}" fill="none" stroke="{GREEN}" stroke-width="4" stroke-dasharray="{DASH}"/>')
    body.append(f'<path d="M{" L".join(f"{num(px(s))} {num(py(c))}" for s, c in zip(show.SHARES, refused))}" fill="none" stroke="{ORANGE}" stroke-width="4"/>')
    for s, c in zip(show.SHARES, changed):
        body.append(square(px(s), py(c)))
    for s, c in zip(show.SHARES, refused):
        body.append(circle(px(s), py(c), 7, ORANGE))
    body.append(circle(250, py(75) - 7, 7, ORANGE))
    body.append(text(266, py(75), "flagged requests still refused"))
    body.append(square(250, py(40) - 7))
    body.append(text(266, py(40), "plain answers changed"))
    at = lambda series, share: series[show.SHARES.index(share)]
    desc = (f"Requests out of 100 against the share of the direction taken out, read after layer {LAYER}. Flagged requests still refused (orange line, filled circles): "
            f"{at(refused, 0)} at 0%, {at(refused, 10)} at 10%, {at(refused, 15)} at 15%, {at(refused, 20)} at 20%, {at(refused, 25)} at 25% and {at(refused, 30)} from 30% on. "
            f"Plain answers changed (green dashed line, hollow squares): {at(changed, 0)} at 0%, {at(changed, 5)} from 5% on.")
    return frame(600, base + 84, "Refusals fall in a narrow band of the share taken out, while ordinary answers hardly move", desc, body)


# --- figure 4: where to read the direction ----------------------------------------------------------------------------------------

def fig_where_to_read() -> str:
    r = show.run()
    top, step = 100, 44
    body = [text(150, 34, "degrees off", bold=True), text(440, 34, "still", "middle", bold=True), text(540, 34, "answers", "middle", bold=True),
            text(440, 60, "refused", "middle", bold=True), text(540, 60, "changed", "middle", bold=True)]
    rows = []
    for layer in range(LAYERS + 1):
        y = top + layer * step
        length = r["angle"][layer] * 4.5
        refused, changed = r["refused"][layer][-1], r["changed"][layer][-1]
        body.append(text(30, y + 8, f"layer {layer}"))
        body.append(f'<rect x="150" y="{y - 16}" width="{num(length)}" height="30" fill="{PANEL}" stroke="{INK}" stroke-width="2"/>')
        body.append(text(150 + length + 10, y + 8, f"{r['angle'][layer]:.1f}", bold=True))
        body.append(text(440, y + 8, str(refused), "middle", bold=True))
        body.append(text(540, y + 8, str(changed), "middle", bold=True))
        rows.append(f"layer {layer}: {r['angle'][layer]:.1f} degrees off, {refused} refused, {changed} changed")
    desc = (f"For each of the {LAYERS + 1} layers at which the direction can be read, how many degrees the direction found is from the planted one, how many of 100 flagged requests "
            f"are still refused after all of it is taken out, and how many of 100 plain answers change. " + "; ".join(rows) + ".")
    return frame(600, top + (LAYERS + 1) * step - 10, "The later the direction is read, the nearer it is to the planted one and the less it disturbs", desc, body)


# --- figure 5: a second route ---------------------------------------------------------------------------------------------------

def fig_second_route() -> str:
    second = 0.3
    model, sets, base = study(second=second)
    u = direction_at(base, LAYER)
    refused_after = set(np.flatnonzero(forward(model, sets["test_flagged"], remove(u))[0].argmax(1) == 0).tolist())
    refused_before = set(np.flatnonzero(base["test_flagged"][0].argmax(1) == 0).tolist())
    on_second = round(second * len(sets["test_flagged"]))
    inside = sum(k < on_second for k in refused_after)
    top, pitch = 76, 20
    body = [f'<rect x="{left - 5}" y="{top - 5}" width="{10 * pitch + 5}" height="{3 * pitch + 5}" fill="{PANEL}"/>' for left in (30, 330)]
    for title, left, refused in (("Before", 30, refused_before), ("After", 330, refused_after)):
        body.append(text(left, 36, title, bold=True))
        for k in range(len(sets["test_flagged"])):
            x, y = left + (k % 10) * pitch, top + (k // 10) * pitch
            fill = ORANGE if k in refused else "none"
            body.append(f'<rect x="{x}" y="{y}" width="16" height="16" fill="{fill}" stroke="{ORANGE if k in refused else INK}" stroke-width="2"/>')
        body.append(text(left, top + 10 * pitch + 32, f"{len(refused)} refused", bold=True))
    body.append(text(300, top + 10 * pitch + 72, "shaded rows: the second route", "middle", GREY))
    body.append(text(300, top + 10 * pitch + 102, "filled: refused. hollow: answered", "middle", GREY))
    desc = (f"Two grids of the 100 held-out flagged requests when {on_second} of them use a second route (the shaded top three rows), each square one request. Before the edit "
            f"{len(refused_before)} are refused (filled). After the direction read at layer {LAYER} is taken out, {len(refused_after)} are refused, {inside} of them among the {on_second} "
            f"second-route requests.")
    return frame(600, top + 10 * pitch + 122, "One direction does not remove a refusal that travels along two", desc, body)


FIGURES = {
    "the-map.svg": fig_the_map,
    "weight-edit.svg": fig_weight_edit,
    "share-removed.svg": fig_share_removed,
    "where-to-read.svg": fig_where_to_read,
    "second-route.svg": fig_second_route,
}


def main() -> int:
    stale = []
    for name, draw in FIGURES.items():
        path, svg = FIGS / name, draw()
        if "--check" in sys.argv:
            if not path.exists() or path.read_text(encoding="utf-8") != svg:
                stale.append(name)
        else:
            FIGS.mkdir(exist_ok=True)
            path.write_text(svg, encoding="utf-8")
    if "--check" in sys.argv:
        print("figures differ from refusal_toy.py: " + ", ".join(stale) if stale else "figures match refusal_toy.py")
        return 1 if stale else 0
    print("wrote " + ", ".join(FIGURES))
    return 0


if __name__ == "__main__":
    sys.exit(main())