"""Draw the four figures of lesson 29 from capacity_plan.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_capacity_plan.py runs the same check)

All four 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. The two limits of a GPU differ by pattern and by label, not by colour alone.
"""
import math
import sys
from pathlib import Path
from xml.sax.saxutils import escape

import capacity_plan as cp

FIGS = Path(__file__).resolve().parent.parent / "figs"
INK, GREY, ORANGE, GREEN, FILL, PANEL = "#0a1626", "#44566b", "#b36200", "#00855F", "#d2dae4", "#ebf0f6"
DASH = "10 8"


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_capacity_plan.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]) -> 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>", '  <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 seconds_text(x: float) -> str:
    return f"{math.floor(x * 100 + 0.5) / 100:.2f}".rstrip("0").rstrip(".")


# --- figure 2: chats in progress through the load test ------------------------------------------------------------------------

def fig_in_progress() -> str:
    load = cp.opening_load_test()
    n = load["plan"]["without"]                                                    # the GPUs the arithmetic alone gives run about as full as the plan assumes
    run = load["runs"][n]
    curve = cp.in_progress_curve(run["spans"], cp.LOAD_SECONDS)
    people = cp.USERS * cp.ACTIVE // 100
    x0, x1, top, base, ymax = 80, 580, 76, 330, 1000
    px = lambda t: x0 + t / cp.LOAD_SECONDS * (x1 - x0)
    py = lambda chats: base - chats / ymax * (base - top)
    body = [text(20, 40, "chats in progress", color=GREY)]
    for tick in (0, 500, 1000):
        body.append(f'<path d="M{x0} {num(py(tick))}H{x1}" stroke="{FILL}" stroke-width="2"/>')
        body.append(text(x0 - 8, py(tick) + 7, cp.commas(tick), "end", GREY, size=20))
    body.append(f'<path d="M{x0} {num(py(cp.USERS))}H{x1}" stroke="{ORANGE}" stroke-width="4" stroke-dasharray="{DASH}"/>')
    body.append(text(x0 + 8, py(cp.USERS) - 10, f"{cp.commas(cp.USERS)} users", bold=True))
    body.append(f'<path d="M{x0} {num(py(people))}H{x1}" stroke="{GREY}" stroke-width="3" stroke-dasharray="{DASH}"/>')
    body.append(text(x0 + 8, py(people) - 10, f"{cp.commas(people)} active in the hour"))
    points = " L".join(f"{num(px(t))} {num(py(curve[t]))}" for t in range(0, cp.LOAD_SECONDS, 3))
    body.append(f'<path d="M{points}" fill="none" stroke="{INK}" stroke-width="3"/>')
    body.append(text(x0 + 8, py(run["in_progress"]) - 28, f"about {round(run['in_progress'])} chats in progress", bold=True))
    body.append(f'<path d="M{x0} {base}H{x1}" stroke="{GREY}" stroke-width="2"/>')
    for minutes in (0, 5, 10):
        body.append(text(px(minutes * 60), base + 32, str(minutes), "middle", GREY, size=20))
    body.append(text((x0 + x1) / 2, base + 62, "minutes of the load test", "middle", GREY))
    desc = (f"Chats in progress through {cp.LOAD_SECONDS // 60} minutes of the load test on {n} GPUs, from 0 to {cp.commas(ymax)}. Two dashed lines mark {cp.commas(cp.USERS)} users and "
            f"{cp.commas(people)} users active in the hour. The solid line stays low, near {round(run['in_progress'])} chats in progress on average.")
    return frame(600, 400, f"Chats in progress: about {round(run['in_progress'])}, for {cp.commas(cp.USERS)} users", desc, body)


# --- figure 1: the two limits of one GPU ------------------------------------------------------------------------------------

def fig_limits() -> str:
    rows = cp.step_table("usual")
    speed, memory = cp.speed_limit(rows, cp.TARGET_MS), cp.memory_limit("usual")
    x0, x1, top, base, xmax, ymax = 84, 580, 70, 320, 250, 200
    px = lambda chats: x0 + chats / xmax * (x1 - x0)
    py = lambda ms: base - ms / ymax * (base - top)
    body = [text(20, 40, "p99 time between words, ms", color=GREY)]
    for tick in (0, 50, 100, 150, 200):
        body.append(f'<path d="M{x0} {num(py(tick))}H{x1}" stroke="{FILL}" stroke-width="2"/>')
        body.append(text(x0 - 8, py(tick) + 7, str(tick), "end", GREY, size=20))
    points = [(0, 0)] + [(r[0], r[3]) for r in rows]
    curve = "M" + " L".join(f"{num(px(c))} {num(py(ms))}" for c, ms in points)
    body.append(f'<path d="{curve}" fill="none" stroke="{INK}" stroke-width="3"/>')
    body.extend(f'<circle cx="{num(px(r[0]))}" cy="{num(py(r[3]))}" r="4" fill="{INK}"/>' for r in rows)
    body.append(f'<path d="M{x0} {num(py(cp.TARGET_MS))}H{x1}" stroke="{ORANGE}" stroke-width="4" stroke-dasharray="{DASH}"/>')
    body.append(text(px(60), py(cp.TARGET_MS) + 30, f"target {cp.TARGET_MS} ms", bold=True))
    body.append(f'<path d="M{num(px(speed))} {base}V{top}" stroke="{GREEN}" stroke-width="4"/>')
    body.append(f'<circle cx="{num(px(speed))}" cy="{num(py(cp.TARGET_MS))}" r="9" fill="{GREEN}" stroke="{INK}" stroke-width="2.5"/>')
    body.append(text(px(speed) + 12, py(190), f"speed limit {math.floor(speed)} chats", bold=True))
    body.append(f'<path d="M{num(px(memory))} {base}V{top}" stroke="{GREY}" stroke-width="4" stroke-dasharray="{DASH}"/>')
    body.append(text(px(memory) - 12, py(160), f"memory holds {memory}", "end", bold=True))
    body.append(f'<path d="M{x0} {base}H{x1}" stroke="{GREY}" stroke-width="2"/>')
    for tick in range(0, xmax + 1, 50):
        body.append(text(px(tick), base + 30, str(tick), "middle", GREY, size=20))
    body.append(text((x0 + x1) / 2, base + 62, "chats in progress on one GPU", "middle", GREY))
    desc = (f"The p99 time between words on one GPU against the chats in progress on it, for usual chats. It rises from {rows[0][3]:.0f} ms with {rows[0][0]:.0f} chat to "
            f"{rows[-1][3]:.0f} ms with {rows[-1][0]:.0f}. A dashed orange line marks the {cp.TARGET_MS} ms target, which the curve crosses at {math.floor(speed)} chats: the speed limit, "
            f"a solid green line. A dashed grey line at {memory} chats is the memory limit, where the cache pool is full.")
    return frame(600, 400, f"One GPU: speed limit {math.floor(speed)} chats, memory limit {memory}", desc, body)


# --- figure 3: which limit comes first --------------------------------------------------------------------------------------

def fig_flip() -> str:
    x0, longest = 112, 330
    rows = []
    for mix, m in cp.MIXES.items():
        speed, memory = math.floor(cp.speed_limit(cp.step_table(mix), cp.TARGET_MS)), cp.memory_limit(mix)
        rows.append((m["label"], memory, speed, "speed" if speed < memory else "memory"))
    scale = longest / max(max(memory, speed) for _, memory, speed, _ in rows)                      # one scale for every bar: the longest limit of any kind of chat
    body = ['<defs><pattern id="hatch" width="10" height="10" patternUnits="userSpaceOnUse" patternTransform="rotate(45)">'
            f'<rect width="10" height="10" fill="#ffffff"/><path d="M0 0V10" stroke="{INK}" stroke-width="4"/></pattern></defs>']
    y, facts = 30, []
    for label, memory, speed, binds in rows:
        body.append(text(20, y, label, bold=True))
        for k, (name, value) in enumerate((("memory", memory), ("speed", speed))):
            top = y + 14 + k * 46
            width = max(value * scale, 3)
            fill = PANEL if name == "memory" else "url(#hatch)"
            border = ORANGE if name == binds else INK
            body.append(text(20, top + 26, name, color=GREY))
            body.append(f'<rect x="{x0}" y="{top}" width="{num(width)}" height="36" fill="{fill}" stroke="{border}" stroke-width="{4 if name == binds else 2}"/>')
            body.append(text(x0 + width + 10, top + 26, f"{value}" + (" binds" if name == binds else ""), bold=name == binds))
        facts.append(f"{label}: memory {memory}, speed {speed}, {binds} first")
        y += 162
    desc = ("Chats in progress one GPU can carry at a p99 of " + f"{cp.TARGET_MS} ms, by chat mix. Each mix has two bars, all on one scale: memory (plain) and speed (striped). The smaller one, "
            "outlined in orange, binds. " + "; ".join(facts) + ".")
    return frame(600, y - 24, "Which limit comes first depends on the chats", desc, body)


# --- figure 4: the load test ------------------------------------------------------------------------------------------------

def fig_load_test() -> str:
    runs = cp.opening_load_test()["runs"]
    x0, longest = 100, 300
    panels = (("p99 time between words, ms", "gap", cp.TARGET_MS, 60, lambda v: f"{v:.1f} ms", f"target {cp.TARGET_MS} ms"),
              ("p99 wait before the first word, s", "wait", cp.WAIT_TARGET, 5, lambda v: f"{seconds_text(v)} s", f"target {seconds_text(cp.WAIT_TARGET)} s"))
    body, y, facts = [], 30, []
    for title, key, target, top_value, show, target_text in panels:
        scale = longest / top_value
        body.append(text(20, y, title, bold=True))
        tx = x0 + target * scale
        for k, (count, run) in enumerate(runs.items()):
            top = y + 16 + k * 46
            ok = run[key] <= target
            body.append(text(20, top + 26, f"{count} GPUs", color=GREY))
            body.append(f'<rect x="{x0}" y="{top}" width="{num(run[key] * scale)}" height="36" fill="{PANEL}" stroke="{GREEN if ok else ORANGE}" stroke-width="4"/>')
            body.append(text(max(x0 + run[key] * scale, tx) + 10, top + 26, show(run[key]) + (", passes" if ok else ", misses"), bold=not ok))
        body.append(f'<path d="M{num(tx)} {y + 8}V{y + 16 + 3 * 46 - 8}" stroke="{INK}" stroke-width="2" stroke-dasharray="6 5"/>')
        body.append(text(tx, y + 16 + 3 * 46 + 18, target_text, "middle", GREY))
        y += 16 + 3 * 46 + 52
    for count, run in runs.items():
        facts.append(f"{count} GPUs: {run['gap']:.1f} ms and {seconds_text(run['wait'])} s, {'passes' if run['passes'] else 'misses'}")
    desc = (f"A {cp.LOAD_SECONDS // 60}-minute load test of {len(cp.opening_load_test()['requests']):,} random arrivals, {cp.peak_rate(cp.USERS, cp.ACTIVE, cp.REPLIES):.1f} a second, through "
            f"{' and '.join(str(c) for c in runs)} GPUs. Two bar charts, with a dashed line at each target. " + "; ".join(facts) + ".")
    return frame(600, y - 14, "The load test by number of GPUs", desc, body)


FIGURES = {
    "chats-in-progress.svg": fig_in_progress,
    "two-limits.svg": fig_limits,
    "which-limit-binds.svg": fig_flip,
    "load-test.svg": fig_load_test,
}


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 capacity_plan.py: " + ", ".join(stale) if stale else "figures match capacity_plan.py")
        return 1 if stale else 0
    print("wrote " + ", ".join(FIGURES))
    return 0


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