"""Draws the three figures that show numbers from the filter: scores.svg, sweep.svg and rare.svg.

The two other figures (flow.svg, outcomes.svg) are drawn by hand. Run `python3 make_figures.py` to
write the three files into ../figs, or `python3 make_figures.py --check` to fail if any is out of date.
Standard library only, plus the filter in safety_filter.py.
"""
from __future__ import annotations

import sys
from pathlib import Path

import safety_filter as sf

FIGS = Path(__file__).resolve().parent.parent / "figs"
INK, GREY, ORANGE, GREEN, LIGHT = "#0a1626", "#44566b", "#b36200", "#00855F", "#d2dae4"

HEAD = (
    '<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 {w} {h}" width="{w}" height="{h}" role="img" '
    'font-family="system-ui, sans-serif" font-size="22" fill="{ink}">\n'
    "  <title>{title}</title>\n  <desc>{desc}</desc>\n"
    '  <rect width="100%" height="100%" fill="#ffffff"/>\n'
)


def head(w: int, h: int, title: str, desc: str) -> str:
    return HEAD.format(w=w, h=h, title=title, desc=desc, ink=INK)


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


def pct(x: float) -> str:
    return f"{x * 100:.2f}".rstrip("0").rstrip(".") + "%"


# --- scores.svg: how many test messages got each score ------------------------------------------------

def bins(labels, score_list):
    """Ten bins of 0.1. Returns (harmful per bin, harmless per bin). Scores have 3 decimals, so work in thousandths."""
    out = ([0] * 10, [0] * 10)
    for y, s in zip(labels, score_list):
        out[1 - y][min(round(s * 1000) // 100, 9)] += 1
    return out


def scores_svg(labels, score_list) -> str:
    harmful, harmless = bins(labels, score_list)
    x0, bin_w, base, unit = 60, 52, 150, 3.2
    cut = x0 + 0.5 * 10 * bin_w
    caught, missed, false_alarms, _ = sf.counts(labels, score_list, 50)
    out = [head(600, 430, "How many test messages got each score",
                f"A mirrored bar chart of the {len(labels)} test messages by score. Harmful messages are orange bars "
                f"going up, harmless messages are green bars going down. Harmless messages crowd the low scores "
                f"({harmless[0]} of them score under 0.1). Harmful messages crowd the top ({harmful[9]} score above 0.9). "
                f"A dashed line marks the threshold 0.50. {missed} harmful messages sit left of it and {false_alarms} "
                f"harmless message sits right of it. Those are the mistakes.")]
    for i in range(0, 11, 2):
        x = x0 + i * bin_w
        out.append(f'  <path d="M{x} 40V350" stroke="{LIGHT}" stroke-width="2"/>')
    out.append(f'  <path d="M{x0} {base}H{x0 + 10 * bin_w}" stroke="{GREY}" stroke-width="2"/>')
    bars, texts = [], []
    for i in range(10):
        x = x0 + i * bin_w + 4
        left = i < 5
        if harmful[i]:
            h = harmful[i] * unit
            stroke = f' stroke="{INK}" stroke-width="3"' if left else ""
            bars.append(f'  <rect x="{x}" y="{num(base - h)}" width="{bin_w - 8}" height="{num(h)}" fill="{ORANGE}"{stroke}/>')
            texts.append(f'    <text x="{x + (bin_w - 8) / 2}" y="{num(base - h - 8)}">{harmful[i]}</text>')
        if harmless[i]:
            h = harmless[i] * unit
            stroke = f' stroke="{INK}" stroke-width="3"' if not left else ""
            bars.append(f'  <rect x="{x}" y="{base}" width="{bin_w - 8}" height="{num(h)}" fill="{GREEN}"{stroke}/>')
            texts.append(f'    <text x="{x + (bin_w - 8) / 2}" y="{num(base + h + 24)}">{harmless[i]}</text>')
    out += bars
    out.append(f'  <path d="M{cut} 44V350" stroke="{INK}" stroke-width="3" stroke-dasharray="10 6"/>')
    out.append('  <g text-anchor="middle">')
    out += texts
    out.append(f'    <text x="{cut}" y="26" font-weight="bold">threshold 0.50</text>')
    out.append("  </g>")
    out.append(f'  <text x="{cut - 8}" y="68" text-anchor="end">&#8592; passed</text>')
    out.append(f'  <text x="{cut + 8}" y="68">blocked &#8594;</text>')
    out.append(f'  <rect x="{cut + 20}" y="92" width="16" height="16" fill="{ORANGE}" stroke="{INK}" stroke-width="2"/>')
    out.append(f'  <text x="{cut + 44}" y="108" font-weight="bold">Harmful ({sum(harmful)})</text>')
    out.append(f'  <circle cx="{cut + 16}" cy="236" r="9" fill="{GREEN}" stroke="{INK}" stroke-width="2"/>')
    out.append(f'  <text x="{cut + 34}" y="243" font-weight="bold">Harmless ({sum(harmless)})</text>')
    out.append(f'  <text x="{x0 + 3.5 * bin_w}" y="{base - 52}" text-anchor="middle">{missed} missed</text>')
    out.append(f'  <text x="{x0 + 6 * bin_w}" y="{base + 56}">{false_alarms} false alarm</text>')
    out.append('  <g font-size="20" text-anchor="middle">')
    for i, label in enumerate(("0", "0.2", "0.4", "0.6", "0.8", "1")):
        out.append(f'    <text x="{x0 + 2 * i * bin_w}" y="380">{label}</text>')
    out.append("  </g>")
    out.append(f'  <text x="{x0 + 5 * bin_w}" y="416" text-anchor="middle">Score the filter gave</text>')
    out.append("</svg>")
    return "\n".join(out) + "\n"


# --- sweep.svg: the two error rates at every threshold ------------------------------------------------

def sweep_svg(labels, score_list) -> str:
    table = sf.sweep(labels, score_list)
    x0, x1, y0, y1, top = 66, 520, 280, 20, 70  # plot box; rates 0 to `top` percent
    px = lambda pct_: x0 + pct_ / 100 * (x1 - x0)  # noqa: E731  (threshold in hundredths -> x)
    py = lambda r: y0 - r * 100 / top * (y0 - y1)  # noqa: E731  (rate 0..1 -> y)
    miss = [(px(p), py(sf.rates(table[p])[0])) for p in range(1, 100)]
    false = [(px(p), py(sf.rates(table[p])[1])) for p in range(1, 100)]
    line = lambda pts: " ".join(f"{num(x)},{num(y)}" for x, y in pts)  # noqa: E731
    m50, f50, _ = sf.rates(table[50])
    cross = sf.crossing(table)
    out = [head(540, 350, "The two error rates at every threshold",
                "A line chart with the threshold from 0 to 1 along the bottom and a rate from 0 to 70 percent up the side. "
                "The miss rate, an orange solid line, is zero at low thresholds and climbs as the threshold rises. "
                "The false-alarm rate, a green dashed line, starts high and falls to zero. "
                f"The lines cross at about threshold {cross / 100:.2f}. At 0.50 the miss rate is {pct(m50)} and the "
                f"false-alarm rate is {pct(f50)}.")]
    for r in (0, 20, 40, 60):
        out.append(f'  <path d="M{x0} {num(py(r / 100))}H{x1}" stroke="{LIGHT}" stroke-width="2"/>')
    out.append('  <g font-size="20" text-anchor="end">')
    for r in (0, 20, 40, 60):
        out.append(f'    <text x="{x0 - 8}" y="{num(py(r / 100) + 7)}">{r}%</text>')
    out.append("  </g>")
    out.append('  <g font-size="20" text-anchor="middle">')
    for t in (0, 25, 50, 75, 100):
        out.append(f'    <text x="{num(px(t))}" y="310">{t / 100:g}</text>')
    out.append("  </g>")
    out.append(f'  <path d="M{num(px(50))} {y1}V{y0}" stroke="{GREY}" stroke-width="2" stroke-dasharray="4 5"/>')
    out.append(f'  <polyline points="{line(false)}" fill="none" stroke="{GREEN}" stroke-width="4" stroke-dasharray="12 7"/>')
    out.append(f'  <polyline points="{line(miss)}" fill="none" stroke="{ORANGE}" stroke-width="4"/>')
    out.append(f'  <circle cx="{num(px(50))}" cy="{num(py(m50))}" r="7" fill="{ORANGE}" stroke="{INK}" stroke-width="2"/>')
    out.append(f'  <circle cx="{num(px(50))}" cy="{num(py(f50))}" r="7" fill="{GREEN}" stroke="{INK}" stroke-width="2"/>')
    halo = f'stroke="#ffffff" stroke-width="7" paint-order="stroke"'
    out.append(f'  <text x="{num(px(50))}" y="{num(py(m50) - 48)}" text-anchor="middle" {halo}>{pct(m50)} missed</text>')
    out.append(f'  <text x="{num(px(50) + 14)}" y="{num(py(f50) - 14)}" {halo}>{pct(f50)} false alarms</text>')
    out.append(f'  <text x="{x0 + 14}" y="44" font-weight="bold">False-alarm rate</text>')
    out.append(f'  <text x="{num(px(80))}" y="{num(py(sf.rates(table[80])[0]) - 16)}" text-anchor="end" font-weight="bold">Miss rate</text>')
    out.append(f'  <text x="{num((x0 + x1) / 2)}" y="340" text-anchor="middle">Threshold</text>')
    out.append("</svg>")
    return "\n".join(out) + "\n"


# --- rare.svg: who is in the blocked pile ---------------------------------------------------------------

def rare_svg(labels, score_list) -> str:
    c = sf.counts(labels, score_list, 50)
    caught_t, _, false_t, _ = sf.at_traffic(c)
    rows = [("On the test messages", c[0], c[2], 82, ""), ("In 1 in 100 traffic", caught_t, false_t, 250, "about ")]
    width, x0 = 460, 40
    out = [head(540, 376, "Who is in the blocked pile",
                f"Two bars show everything blocked at threshold 0.50. On the test messages {c[0] + c[2]} are blocked: "
                f"{c[0]} harmful and {c[2]} harmless. In traffic where 1 message in 100 is harmful, "
                f"about {caught_t + false_t} of every 10,000 are blocked: about {caught_t} harmful and about {false_t} harmless.")]
    for name, harm, ok, y, about in rows:
        total = harm + ok
        w_harm = width * harm / total
        out.append(f'  <text x="{x0}" y="{y - 14}" font-weight="bold">{name}: {about}{total} blocked</text>')
        out.append(f'  <rect x="{x0}" y="{y}" width="{num(w_harm)}" height="56" fill="{ORANGE}" stroke="{INK}" stroke-width="2"/>')
        out.append(f'  <rect x="{num(x0 + w_harm)}" y="{y}" width="{num(width - w_harm)}" height="56" fill="{GREEN}" stroke="{INK}" stroke-width="2"/>')
        out.append(f'  <text x="{x0}" y="{y + 88}">{about}{harm} harmful</text>')
        out.append(f'  <text x="{x0}" y="{y + 116}">{round(100 * harm / total)}%</text>')
        out.append(f'  <text x="{x0 + width}" y="{y + 88}" text-anchor="end">{about}{ok} harmless</text>')
        out.append(f'  <text x="{x0 + width}" y="{y + 116}" text-anchor="end">{round(100 * ok / total)}%</text>')
    out.append("</svg>")
    return "\n".join(out) + "\n"


def figures() -> dict[str, str]:
    r = sf.build()
    return {"scores.svg": scores_svg(r["labels"], r["scores"]),
            "sweep.svg": sweep_svg(r["labels"], r["scores"]),
            "rare.svg": rare_svg(r["labels"], r["scores"])}


if __name__ == "__main__":
    wanted = figures()
    if "--check" in sys.argv:
        stale = [n for n, svg in wanted.items() if not (FIGS / n).exists() or (FIGS / n).read_text() != svg]
        print("figures out of date: " + ", ".join(stale) if stale else "figures up to date")
        sys.exit(1 if stale else 0)
    for name, svg in wanted.items():
        (FIGS / name).write_text(svg)
        print("wrote", FIGS / name)