"""Draw the four figures of lesson 33 from pd_sim.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_lesson.py runs the same check)

All four figures are 540 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 math
import sys
from functools import lru_cache
from pathlib import Path
from xml.sax.saxutils import escape

import pd_sim as sim

FIGS = Path(__file__).resolve().parent.parent / "figs"
INK, GREY, ORANGE, GREEN, FILL, PANEL = "#0a1626", "#44566b", "#b36200", "#00855F", "#d2dae4", "#ebf0f6"
DASH = "10 8"
PAD = 14                                            # room between the axis and the first mark
RPS, PROMPT, REPLY, LINK = 4, 2000, 200, 25         # the lesson's default run
RPS_GRID = tuple(range(1, 11))
PROMPT_GRID = (100, 200, 400, 600, 800, 1000, 1500, 2000, 3000, 4000, 6000)
LINK_GRID = (0.1, 0.125, 0.15, 0.2, 0.25, 0.29, 0.3, 0.4, 0.5, 0.75, 1, 1.25, 1.5, 2, 2.5, 3, 5, 12.5, 25, 100)   # includes the five speeds of sim.LINKS_GBPS


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_lesson.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 tick(x: float, y: float, label: str, anchor: str = "middle") -> str:
    return text(x, y, label, anchor, GREY, size=20)


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 triangle(x: float, y: float) -> str:
    return f'<path d="M{num(x)} {num(y - 9)}L{num(x + 9)} {num(y + 7)}L{num(x - 9)} {num(y + 7)}z" fill="{ORANGE}" stroke="#ffffff" stroke-width="1"/>'


def block(x0: float, x1: float, y: float, height: float, stroke: str, dash: str = "") -> str:
    return (f'<rect x="{num(x0)}" y="{num(y)}" width="{num(x1 - x0)}" height="{num(height)}" fill="{FILL}" stroke="{stroke}" stroke-width="3"'
            + (f' stroke-dasharray="{dash}"' if dash else "") + "/>")


def ms(seconds: float) -> str:
    return f"{seconds * 1000:.1f}"


def path(points: list[tuple[float, float]], color: str, dash: str = "") -> str:
    return f'<path d="M{" L".join(f"{num(x)} {num(y)}" for x, y in points)}" fill="none" stroke="{color}" stroke-width="4"' + (f' stroke-dasharray="{dash}"' if dash else "") + "/>"


def key(x: float, y: float, label: str, kind: str) -> list[str]:
    """One line of a key: a short sample of the mark, then its name. kind is shared (orange line, filled circle), split (green dashed line, hollow square) or rule (grey dash)."""
    if kind == "rule":
        return [f'<path d="M{num(x - 18)} {num(y - 7)}H{num(x + 18)}" stroke="{GREY}" stroke-width="3" stroke-dasharray="{DASH}"/>', text(x + 30, y, label)]
    line = path([(x - 18, y - 7), (x + 18, y - 7)], ORANGE if kind == "shared" else GREEN, "" if kind == "shared" else DASH)
    return [line, circle(x, y - 7, 7, ORANGE) if kind == "shared" else square(x, y - 7), text(x + 30, y, label)]


# --- the numbers: each is computed once from the simulation ------------------------------------------------------------------------

@lru_cache(maxsize=None)
def runs(rps: float, prompt: int, reply: int, link: float) -> dict:
    return sim.compare(rps, prompt, reply, link)


@lru_cache(maxsize=None)
def crossover_gbps(rps: float = RPS, prompt: int = PROMPT, reply: int = REPLY) -> float:
    """The slowest link at which the split's goodput is still above the shared layout's, found by halving the gap between a link that loses and one that wins."""
    lose, win = 0.1, 100.0
    assert runs(rps, prompt, reply, lose)["split"]["goodput"] <= runs(rps, prompt, reply, win)["shared"]["goodput"] < runs(rps, prompt, reply, win)["split"]["goodput"]
    for _ in range(40):
        middle = math.sqrt(lose * win)
        if runs(rps, prompt, reply, middle)["split"]["goodput"] > runs(rps, prompt, reply, middle)["shared"]["goodput"]:
            win = middle
        else:
            lose = middle
    return win


@lru_cache(maxsize=None)
def stall_scene() -> dict:
    """Figure 1: reply A is in progress when B, a long prompt, arrives; the same two requests on one shared GPU and on a split pair."""
    a_read, b_read, arrival = 100 * sim.READ_S_PER_TOKEN, PROMPT * sim.READ_S_PER_TOKEN, 0.045
    first, gaps = sim.serve([(0.0, a_read, 40, 0), (arrival, b_read, 5, 1)])
    requests = [sim.Request(0.0, 100, 41), sim.Request(arrival, PROMPT, 6)]
    split_first, split_gaps = sim.split(requests, LINK)

    def words(start: float, steps: list[float]) -> list[float]:
        times = [start]
        for gap in steps:
            times.append(times[-1] + gap)
        return times

    return {"shared_words": words(first[0], gaps[0]), "read_end": first[1], "read_start": first[1] - b_read, "split_words": words(split_first[0], split_gaps[0]),
            "arrival": arrival, "reader_start": arrival, "reader_end": arrival + b_read, "landed": split_first[1], "hand_over": sim.handover_seconds(PROMPT, LINK)}


# --- figure 1: one long prompt, on a shared GPU and on a split pair -----------------------------------------------------------------

def fig_stall() -> str:
    s = stall_scene()
    x0, window, eps = 20, 0.24, 1e-9
    px = lambda t: x0 + t / window * 500
    words = [t for t in s["shared_words"] if t <= window]
    before = max(t for t in words if t <= s["read_start"] + eps)
    after = min(t for t in s["shared_words"] if t >= s["read_end"] + sim.STEP_S - eps)
    body = [text(x0, 28, "One shared GPU", bold=True)]
    body.append(block(px(s["read_start"]), px(s["read_end"]), 40, 36, ORANGE))
    body.append(text((px(s["read_start"]) + px(s["read_end"])) / 2, 64, f"reads B: {ms(s['read_end'] - s['read_start'])} ms", "middle"))
    body += [circle(px(t), 58) for t in words if t <= s["read_start"] + eps or t >= after - eps]
    body.append(f'<path d="M{num(px(before))} 102H{num(px(after))}M{num(px(before))} 95V109M{num(px(after))} 95V109" stroke="{INK}" stroke-width="3"/>')
    body.append(text((px(before) + px(after)) / 2, 132, f"{ms(after - before)} ms with no word from A", "middle"))
    body.append(triangle(px(s["read_end"]), 88))
    body.append(text(x0, 168, "Split: reader GPU", bold=True))
    body.append(block(px(s["reader_start"]), px(s["reader_end"]), 180, 36, ORANGE))
    body.append(text((px(s["reader_start"]) + px(s["reader_end"])) / 2, 204, f"reads B: {ms(s['reader_end'] - s['reader_start'])} ms", "middle"))
    body.append(text(x0, 248, "Split: link", bold=True))
    body.append(block(px(s["reader_end"]), px(s["landed"]), 260, 28, GREY, DASH))
    body.append(text(px(s["reader_end"]) - 12, 282, f"B's cache crosses: {ms(s['hand_over'])} ms", "end"))
    body.append(text(x0, 330, "Split: writer GPU", bold=True))
    body.append(text(x0 + 500, 330, f"A: a word every {ms(sim.STEP_S)} ms", "end"))
    body += [circle(px(t), 358) for t in s["split_words"] if t <= window]
    body.append(triangle(px(s["landed"]), 388))
    body.append(f'<path d="M{x0} 412H{x0 + 500}" stroke="{GREY}" stroke-width="2"/>')
    for t in (0, 0.05, 0.10, 0.15, 0.20):
        body.append(tick(px(t), 440, f"{round(t * 1000)}"))
    body.append(text(x0 + 500, 470, "time, in ms", "end", GREY))
    body.append(circle(x0 + 8, 463, 7))
    body.append(text(x0 + 24, 470, "A's word"))
    body.append(triangle(x0 + 150, 463))
    body.append(text(x0 + 168, 470, "B's first word"))
    desc = (f"Time in milliseconds on four rows. One shared GPU: reply A gets a word every {ms(sim.STEP_S)} ms until B's {PROMPT}-token prompt arrives at {ms(s['arrival'])} ms. "
            f"The GPU then reads B from {ms(s['read_start'])} to {ms(s['read_end'])} ms ({ms(s['read_end'] - s['read_start'])} ms), and A gets no word for {ms(after - before)} ms. "
            f"Split: the reader GPU reads B from {ms(s['reader_start'])} to {ms(s['reader_end'])} ms, B's cache crosses the link in {ms(s['hand_over'])} ms and lands at {ms(s['landed'])} ms, "
            f"and the writer GPU gives A a word every {ms(sim.STEP_S)} ms throughout.")
    return frame(540, 486, f"A {PROMPT}-token prompt stops a shared GPU's replies for {ms(after - before)} ms, and does not stop a writer GPU", desc, body)


# --- figures 2 to 4: lines against a setting -----------------------------------------------------------------------------------------

def axes(x0: float, x1: float, top: float, base: float, yticks: list[tuple[float, str]], header: str, xlabel: str) -> list[str]:
    body = [text(20, 36, header, color=GREY)]
    for y, label in yticks:
        body.append(f'<path d="M{x0} {num(y)}H{x1}" stroke="{FILL}" stroke-width="2"/>')
        body.append(tick(x0 - 8, y + 7, label, "end"))
    body.append(f'<path d="M{x0} {base}H{x1}" stroke="{GREY}" stroke-width="2"/>')
    body.append(text((x0 + x1) / 2, base + 64, xlabel, "middle", GREY))
    return body


def fig_gap_by_load() -> str:
    x0, x1, top, base = 84, 520, 60, 300
    px = lambda r: x0 + PAD + (r - RPS_GRID[0]) / (RPS_GRID[-1] - RPS_GRID[0]) * (x1 - x0 - 2 * PAD)
    py = lambda v: base - v / 0.25 * (base - top)
    got = {r: runs(r, PROMPT, REPLY, LINK) for r in RPS_GRID}
    gap = lambda layout: [got[r][layout]["gap_p99"] for r in RPS_GRID]
    body = axes(x0, x1, top, base, [(py(0), "0"), (py(0.1), "100"), (py(0.2), "200")], "p99 gap between words, in ms", "requests per second")
    body.append(f'<path d="M{x0} {num(py(sim.WORD_GAP_TARGET_S))}H{x1}" stroke="{GREY}" stroke-width="3" stroke-dasharray="{DASH}"/>')
    for r in RPS_GRID:
        body.append(tick(px(r), base + 30, str(r)))
    body.append(path([(px(r), py(g)) for r, g in zip(RPS_GRID, gap("shared"))], ORANGE))
    body.append(path([(px(r), py(g)) for r, g in zip(RPS_GRID, gap("split"))], GREEN, DASH))
    body += [circle(px(r), py(g), 7, ORANGE) for r, g in zip(RPS_GRID, gap("shared"))]
    body += [square(px(r), py(g)) for r, g in zip(RPS_GRID, gap("split"))]
    body += key(x0 + 40, top + 12, "shared GPUs", "shared") + key(x0 + 40, top + 40, "split", "split") + key(x0 + 40, top + 68, "100 ms limit", "rule")
    at = lambda layout, r: ms(got[r][layout]["gap_p99"])
    over = min(r for r in RPS_GRID if got[r]["shared"]["gap_p99"] > sim.WORD_GAP_TARGET_S)
    desc = (f"p99 gap between words in ms against requests per second, from 1 to 10, with {PROMPT}-token prompts and {REPLY}-word replies (before the mix of lengths). "
            f"Shared GPUs (orange line, filled circles): {at('shared', 1)} at 1, {at('shared', 3)} at 3, {at('shared', 4)} at 4, {at('shared', 6)} at 6 and {at('shared', 10)} at 10. "
            f"Split (green dashed line, hollow squares): {at('split', 1)} at 1 and {at('split', 10)} at 10. A grey dashed line marks the 100 ms limit.")
    return frame(540, base + 84, f"On shared GPUs the p99 gap passes {sim.WORD_GAP_TARGET_S * 1000:.0f} ms at {over} requests a second; on a split pair it stays at {at('split', over)} ms", desc, body)


def fig_goodput_by_link() -> str:
    x0, x1, top, base = 84, 520, 60, 300
    low, high = math.log10(LINK_GRID[0]), math.log10(LINK_GRID[-1])
    px = lambda g: x0 + PAD + (math.log10(g) - low) / (high - low) * (x1 - x0 - 2 * PAD)
    py = lambda v: base - v * (base - top)
    got = {g: runs(RPS, PROMPT, REPLY, g) for g in LINK_GRID}
    cut = crossover_gbps()
    shared_goodput = got[LINK]["shared"]["goodput"]
    body = axes(x0, x1, top, base, [(py(0), "0"), (py(0.5), "50%"), (py(1), "100%")], "goodput: requests that meet both targets", "link speed, in GB/s (log scale)")
    for g in (0.1, 1, 10, 100):
        body.append(tick(px(g), base + 30, str(g if g < 1 else int(g))))
    body.append(f'<path d="M{num(px(cut))} {top}V{base}" stroke="{GREY}" stroke-width="3" stroke-dasharray="{DASH}"/>')
    body.append(path([(x0, py(shared_goodput)), (x1, py(shared_goodput))], ORANGE))
    body.append(path([(px(g), py(got[g]["split"]["goodput"])) for g in LINK_GRID], GREEN, DASH))
    body += [square(px(g), py(got[g]["split"]["goodput"])) for g in sim.LINKS_GBPS]
    body += key(x0 + 190, top + 80, "split", "split") + key(x0 + 190, top + 114, f"shared GPUs, {shared_goodput:.1%}", "shared") + key(x0 + 190, top + 148, "crossover,", "rule")
    body.append(text(x0 + 220, top + 174, f"about {cut:.1f} GB/s"))                                         # the rule's label runs to a second line so it fits the figure's width
    at = lambda g: f"{got[g]['split']['goodput']:.1%}"
    desc = (f"Goodput against link speed from 0.1 to 100 GB/s, on a log scale, at {RPS} requests a second with {PROMPT}-token prompts and {REPLY}-word replies. "
            f"Split (green dashed line, hollow squares at the five speeds on offer): {at(100)} at 100, {at(25)} at 25, {at(12.5)} at 12.5, {at(1.25)} at 1.25 and {at(0.125)} at 0.125 GB/s. "
            f"Shared GPUs (orange line) stay at {shared_goodput:.1%}. A grey dashed line marks the crossover at about {cut:.1f} GB/s: below it no request meets the 0.5 s target for its first word.")
    return frame(540, base + 84, f"A split pair beats shared GPUs on goodput on links faster than about {cut:.1f} GB/s, and loses below it", desc, body)


def fig_goodput_by_prompt() -> str:
    x0, x1, top, base = 84, 520, 60, 300
    low, high = math.log10(PROMPT_GRID[0]), math.log10(PROMPT_GRID[-1])
    px = lambda p: x0 + PAD + (math.log10(p) - low) / (high - low) * (x1 - x0 - 2 * PAD)
    py = lambda v: base - v * (base - top)
    got = {p: runs(RPS, p, REPLY, LINK) for p in PROMPT_GRID}
    series = lambda layout: [got[p][layout]["goodput"] for p in PROMPT_GRID]
    body = axes(x0, x1, top, base, [(py(0), "0"), (py(0.5), "50%"), (py(1), "100%")], "goodput: requests that meet both targets", "tokens in a typical prompt (log scale)")
    for p in (100, 200, 500, 1000, 2000, 5000):
        body.append(tick(px(p), base + 30, f"{p:,}"))
    body.append(path([(px(p), py(v)) for p, v in zip(PROMPT_GRID, series("shared"))], ORANGE))
    body.append(path([(px(p), py(v)) for p, v in zip(PROMPT_GRID, series("split"))], GREEN, DASH))
    body += [circle(px(p), py(v), 7, ORANGE) for p, v in zip(PROMPT_GRID, series("shared"))]
    body += [square(px(p), py(v)) for p, v in zip(PROMPT_GRID, series("split"))]
    body += key(x0 + 30, top + 160, "shared GPUs", "shared") + key(x0 + 30, top + 194, "split", "split")
    at = lambda layout, p: f"{got[p][layout]['goodput']:.1%}"
    desc = (f"Goodput against the typical prompt length, from {PROMPT_GRID[0]} to {PROMPT_GRID[-1]} tokens on a log scale, at {RPS} requests a second with {REPLY}-word replies and a {LINK} GB/s link. "
            f"Shared GPUs (orange line, filled circles): {at('shared', 200)} at 200, {at('shared', 1000)} at 1,000, {at('shared', 2000)} at 2,000 and {at('shared', 4000)} at 4,000. "
            f"Split (green dashed line, hollow squares): {at('split', 200)} at 200, {at('split', 1000)} at 1,000, {at('split', 2000)} at 2,000, {at('split', 4000)} at 4,000 and {at('split', 6000)} at 6,000.")
    tie = max(p for p in PROMPT_GRID if all(abs(got[q]["split"]["goodput"] - got[q]["shared"]["goodput"]) < 0.01 for q in PROMPT_GRID if q <= p))     # the layouts tie up to here
    return frame(540, base + 84, f"The two layouts nearly tie on prompts up to {tie} tokens, and the split pulls ahead after that", desc, body)


FIGURES = {"stall.svg": fig_stall, "gap-by-load.svg": fig_gap_by_load, "goodput-by-link.svg": fig_goodput_by_link, "goodput-by-prompt.svg": fig_goodput_by_prompt}


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


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