# SPDX-License-Identifier: MIT
"""Build the figures for YS-EN-0001 from results.json and the workload files,
and write them into the note between <!-- fig:name --> markers.

Usage: python build_figures.py   (after queue_sim.py)
"""
from __future__ import annotations

import json
import sys
from pathlib import Path

import numpy as np

HERE = Path(__file__).resolve().parent
sys.path.insert(0, str(HERE.parent))
import svgchart as sc  # noqa: E402
from queue_sim import load_cdf  # noqa: E402

PAGE = HERE.parent.parent / "ys-en-0001-homa-tcp-replacement.html"
R = json.loads((HERE / "results.json").read_text())

# colour follows the entity in every figure
SRPT, LEVELS, FAIR = 1, 2, 3
NAMES = {"srpt": "SRPT", "prio8": "8 levels", "ps": "Fair share", "fifo": "FIFO"}
SLOT = {"srpt": SRPT, "prio8": LEVELS, "ps": FAIR}


def x_fmt(v):
    return f"{v:.1f}×" if v < 100 else f"{v:,.0f}×"


def size_label(b):
    b = float(b)
    if b >= 1e6:
        return f"{b / 1e6:g} MB"
    if b >= 1000:
        return f"{b / 1000:.3g} KB"
    return f"{b:.0f} B"


# ---------------------------------------------------------------- figure 1
def fig_cdf():
    sizes, cum = load_cdf("w4")
    p = np.diff(np.concatenate([[0.0], cum]))
    bytes_cum = np.cumsum(p * sizes) / np.sum(p * sizes)
    grid = np.geomspace(50, 1e6, 41)
    idx = [np.searchsorted(sizes, g, side="right") - 1 for g in grid]  # -1: below the smallest message
    msg = [float(cum[i]) * 100 if i >= 0 else 0.0 for i in idx]
    byt = [float(bytes_cum[i]) * 100 if i >= 0 else 0.0 for i in idx]
    series = [
        {"slot": 1, "name": "Share of messages", "end": "messages", "values": msg},
        {"slot": 2, "name": "Share of bytes", "end": "bytes", "values": byt},
    ]
    fr = sc.Frame(720, 330, "Cumulative share of messages and of bytes by message size, workload W4")
    fr.x_scale(50, 1e6, log=True)
    fr.y_scale(0, 100)
    fr.grid_y([0, 25, 50, 75, 100], fmt=lambda v: f"{v:g}%")
    fr.grid_x([100, 1000, 10_000, 100_000, 1_000_000], fmt=size_label, lines=False)
    fr.baseline()
    fr.vline(1000)
    i1k = np.searchsorted(sizes, 1000, side="right") - 1
    fr.note(fr.x(1000) + 6, fr.mt + 12, f"1,000 B: {cum[i1k] * 100:.0f}% of messages,")
    fr.note(fr.x(1000) + 6, fr.mt + 26, f"{bytes_cum[i1k] * 100:.1f}% of bytes")
    fr.titles("Message size (log scale)", "Cumulative share")
    fr.cursor()
    sc.draw_lines(fr, list(grid), series, end_labels=False)
    fr.note(fr.x(300) - 8, fr.y(30), "messages", anchor="end", cls="ink")
    fr.note(fr.x(60_000) + 8, fr.y(30), "bytes", cls="ink")
    sc.column_hits(fr, list(grid), series, lambda i, x: f"Messages ≤ {size_label(x)}", lambda v: f"{v:.1f}%")
    rows = [[size_label(g), f"{m:.1f}%", f"{b:.1f}%"] for g, m, b in zip(grid[::4], msg[::4], byt[::4])]
    return sc.figure(
        code="FIG. 1 — W4 · FACEBOOK HADOOP · SIZE DISTRIBUTION",
        fid="fig-cdf",
        plot=fr.svg(),
        legend_html=sc.legend(series),
        caption=("<b>Figure 1.</b> Most messages are small; most bytes are in large messages. Workload W4 "
                 "(Facebook Hadoop) from Homa's published tables, truncated at 1 MB. Mean message size "
                 f"{R['workloads']['w4']['mean_size']:,.0f} bytes, median 662 bytes."),
        data_html=sc.table("W4 cumulative distribution", ["Size ≤", "Messages", "Bytes"], rows),
    )


# ---------------------------------------------------------------- figure 2 (diagram)
def fig_exchange():
    W, H = 720, 420
    xa, xb, xr = 90, 270, 500
    o = [f'<svg class="dg" viewBox="0 0 {W} {H}" role="img" aria-label="Homa message exchange: unscheduled data, grants, and a short message overtaking a long one">']
    for x, name, sub in ((xa, "Sender A", "200 KB message"), (xb, "Sender B", "300 B message"), (xr, "Receiver", "schedules its downlink")):
        o.append(f'<rect class="box" x="{x - 80}" y="14" width="160" height="44" rx="4"/>')
        o.append(f'<text class="lbl" x="{x}" y="33" text-anchor="middle">{name}</text>')
        o.append(f'<text x="{x}" y="49" text-anchor="middle">{sub}</text>')
        o.append(f'<line class="rule" x1="{x}" x2="{x}" y1="58" y2="{H - 16}"/>')

    def arrow(x1, y1, x2, y2, cls, label, lx=None, ly=None, anchor="middle"):
        o.append(f'<line class="msg {cls}" x1="{x1}" y1="{y1}" x2="{x2}" y2="{y2}"/>')
        d = 1 if x2 > x1 else -1
        o.append(f'<path class="{cls.replace("s", "f")}" d="M{x2},{y2} l{-9 * d},-4.5 l0,9 Z"/>')
        lx = lx if lx is not None else (x1 + x2) / 2
        ly = min(y1, y2) - 7  # above the higher end, so the text clears the line
        o.append(f'<text x="{lx}" y="{ly}" text-anchor="{anchor}">{label}</text>')

    arrow(xa, 84, xr, 110, "s2", "① DATA · first RTTbytes · unscheduled, sent at once", lx=300, ly=88)
    arrow(xr, 136, xa, 162, "s0", "② GRANT · “send up to byte N, at priority P2”", lx=300, ly=140)
    arrow(xa, 178, xr, 204, "s2", "③ DATA · scheduled · low priority P2", lx=300)
    arrow(xb, 236, xr, 262, "s1", "④ DATA · all 300 B unscheduled · top priority P7", lx=390)
    o.append(f'<rect class="f1" x="{xr - 4}" y="266" width="8" height="8" rx="2"/>')
    o.append(f'<text class="lbl" x="{xr + 14}" y="274">⑤ B done first: it</text>')
    o.append(f'<text class="lbl" x="{xr + 14}" y="290">overtook A’s packets</text>')
    o.append(f'<text class="lbl" x="{xr + 14}" y="306">queued at the switch</text>')
    arrow(xr, 318, xa, 344, "s0", "⑥ GRANT · A resumes", lx=300)
    arrow(xa, 362, xr, 388, "s2", "⑦ DATA · scheduled …", lx=300)
    o.append(f'<text x="16" y="{H - 4}">Time runs downward. Higher P = higher switch priority (8 levels, P7 highest, as in the 2018 paper).</text>')
    o.append("</svg>")
    return sc.figure(
        code="FIG. 2 — HOMA MESSAGE EXCHANGE",
        fid="fig-exchange",
        plot="\n".join(o),
        caption=("<b>Figure 2.</b> A Homa exchange, simplified. Sender A blasts its first RTTbytes without "
                 "asking, then waits for grants. The receiver grants A at a low priority because A has many "
                 "bytes left. Sender B's 300-byte message fits inside the unscheduled window, goes out at "
                 "the top priority, and finishes first. Simplified from the design in [1]. In HomaModule "
                 "since September 2026, messages that need grants skip step ① and announce themselves "
                 "instead [6]."),
    )


# ---------------------------------------------------------------- figure 3 (diagram)
def fig_gantt():
    W, H = 720, 250
    x0, x1 = 150, 690
    u = (x1 - x0) / 10
    X = lambda t: x0 + t * u
    o = [f'<svg class="dg" viewBox="0 0 {W} {H}" role="img" aria-label="Two messages on one link under fair sharing and under shortest-first">']
    lanes = [("Fair share", 40), ("Shortest first", 130)]
    for name, y in lanes:
        o.append(f'<text class="lbl" x="{x0 - 14}" y="{y + 26}" text-anchor="end">{name}</text>')
        o.append(f'<rect class="box" x="{x0}" y="{y}" width="{x1 - x0}" height="44" rx="3"/>')

    def bar(t0, t1, y, h, cls):
        o.append(f'<rect class="{cls}" x="{X(t0) + 1:.1f}" y="{y}" width="{X(t1) - X(t0) - 2:.1f}" height="{h}" rx="2"/>')

    y = 40
    bar(0, 1, y + 2, 40, "f2")
    bar(1, 5, y + 2, 19, "f2")
    bar(1, 5, y + 23, 19, "f1")
    bar(5, 10, y + 2, 40, "f2")
    o.append(f'<text class="lbl" x="{X(3):.1f}" y="{y + 60}" text-anchor="middle">B done at t = 5 · took 4 · slowdown 2×</text>')
    y = 130
    bar(0, 1, y + 2, 40, "f2")
    bar(1, 3, y + 2, 40, "f1")
    bar(3, 10, y + 2, 40, "f2")
    o.append(f'<text class="lbl" x="{X(2):.1f}" y="{y + 60}" text-anchor="middle">B done at t = 3 · took 2 · slowdown 1×</text>')
    for t in range(11):
        o.append(f'<text x="{X(t):.1f}" y="{H - 12}" text-anchor="middle">{t}</text>')
    o.append(f'<text x="{x0 - 14}" y="{H - 12}" text-anchor="end">time</text>')
    for yy in (40, 130):
        o.append(f'<text x="{X(10) - 6:.1f}" y="{yy - 6}" text-anchor="end">A done at t = 10</text>')
    o.append(f'<text x="{X(1):.1f}" y="30" text-anchor="middle">B arrives ↓</text>')
    o.append("</svg>")
    leg = '<div class="legend"><span class="k2"><i></i>A · long message, 8 units of work</span><span class="k1"><i></i>B · short message, 2 units</span></div>'
    return sc.figure(
        code="FIG. 3 — ORDER, NOT SPEED",
        fid="fig-gantt",
        plot="\n".join(o),
        legend_html=leg,
        caption=("<b>Figure 3.</b> One link, two messages. Under fair sharing B gets half the link while A is "
                 "active. Under shortest-first B pre-empts A. The link does the same total work either way, "
                 "so A finishes at t = 10 in both lanes; B finishes in half the time."),
    )


# ---------------------------------------------------------------- figure 4
def fig_bysize():
    keys = ["ps", "prio8", "srpt"]
    sizes = [r["size_max"] for r in R["by_size"]["srpt"]]
    xs = list(range(1, 11))
    series = [{"slot": SLOT[k], "name": NAMES[k], "values": [r["p99"] for r in R["by_size"][k]]} for k in keys]
    fr = sc.Frame(720, 340, "99th-percentile slowdown by message-size decile at 80% load, W4")
    fr.x_scale(1, 10)
    fr.y_scale(0.8, 100, log=True)
    fr.grid_y([1, 3, 10, 30, 100], fmt=lambda v: f"{v:g}×")
    fr.x_labels([(i, f"{i * 10}%") for i in xs])
    fr.baseline()
    fr.titles("Message size percentile (deciles, smallest → largest)", "P99 slowdown (log)")
    fr.cursor()
    sc.draw_lines(fr, xs, series)
    prev = [0] + sizes[:-1]
    sc.column_hits(fr, xs, series,
                   lambda i, x: f"Decile {x}: {size_label(prev[i] + 1)}–{size_label(sizes[i])}", x_fmt)
    rows = []
    for i in range(10):
        rows.append([f"{(i + 1) * 10}%", size_label(sizes[i])] +
                    [x_fmt(R["by_size"][k][i]["p99"]) for k in ("ps", "prio8", "srpt", "fifo")])
    return sc.figure(
        code="FIG. 4 — P99 SLOWDOWN BY MESSAGE SIZE · W4 · 80% LOAD",
        fid="fig-bysize",
        plot=fr.svg(),
        legend_html=sc.legend(series),
        caption=("<b>Figure 4.</b> 99th-percentile slowdown for each tenth of messages, smallest to "
                 "largest. Fair sharing is flat at about 21×. Under SRPT the smallest 60% have a P99 slowdown of 1.0, "
                 "and its largest decile still beats fair sharing. First come, first served (in the table) "
                 f"reaches {R['by_size']['fifo'][0]['p99']:,.0f}× for the smallest decile. "
                 f"Simulated, {R['n']:,} messages per policy."),
        data_html=sc.table("P99 slowdown by size decile, W4, 80% load",
                           ["Decile", "Up to", "Fair share", "8 levels", "SRPT", "FIFO"], rows, numeric_from=1),
        wide=False,
    )


# ---------------------------------------------------------------- figure 5
def fig_byload():
    keys = ["ps", "prio8", "srpt"]
    loads = sorted(R["by_load"], key=float)
    xs = [float(l) * 100 for l in loads]
    series = [{"slot": SLOT[k], "name": NAMES[k], "values": [R["by_load"][l][k]["p99"] for l in loads]} for k in keys]
    fr = sc.Frame(720, 320, "99th-percentile slowdown versus network load, W4")
    fr.x_scale(30, 90)
    fr.y_scale(1, 100, log=True)
    fr.grid_y([1, 3, 10, 30, 100], fmt=lambda v: f"{v:g}×")
    fr.grid_x(xs, fmt=lambda v: f"{v:.0f}%", lines=False)
    fr.baseline()
    fr.titles("Network load", "P99 slowdown, all messages (log)")
    fr.cursor()
    sc.draw_lines(fr, xs, series)
    sc.column_hits(fr, xs, series, lambda i, x: f"Load {x:.0f}%", x_fmt)
    rows = [[f"{float(l) * 100:.0f}%"] + [x_fmt(R["by_load"][l][k]["p99"]) for k in ("ps", "prio8", "srpt")] +
            [f'{R["by_load"][l]["ps"]["mean"]:.2f}×', f"{1 / (1 - float(l)):.2f}×"] for l in loads]
    return sc.figure(
        code="FIG. 5 — P99 SLOWDOWN VS LOAD · W4",
        fid="fig-byload",
        plot=fr.svg(),
        legend_html=sc.legend(series),
        caption=("<b>Figure 5.</b> Tail slowdown across all message sizes as load rises. The table view also "
                 "shows the simulator's check: mean fair-share slowdown against the M/G/1 processor-sharing "
                 "result 1/(1 − ρ)."),
        data_html=sc.table("P99 slowdown vs load, W4",
                           ["Load", "Fair share", "8 levels", "SRPT", "Fair share mean (sim)", "1/(1 − ρ)"], rows),
    )


# ---------------------------------------------------------------- figure 6
def fig_levels():
    order = [("ps", "1 level"), ("prio2", "2"), ("prio4", "4"), ("prio8", "8"), ("prio16", "16"), ("srpt", "SRPT")]
    vals = [R["levels"][k]["short"]["p99"] for k, _ in order]
    slots = [FAIR, LEVELS, LEVELS, LEVELS, LEVELS, SRPT]
    fr = sc.Frame(720, 300, "99th-percentile slowdown of the shortest half of messages by number of priority levels",
                  margins=(58, 30, 30, 46))
    fr.x_scale(0, 1)
    fr.y_scale(0, 25)
    fr.grid_y([0, 5, 10, 15, 20, 25], fmt=lambda v: f"{v:g}×")
    fr.baseline()
    fr.titles("Priority levels (1 level = fair sharing; SRPT = unlimited)", "P99 slowdown, shortest 50%")
    sc.columns(fr, [c for _, c in order], vals, slots=slots, fmt=lambda v: f"{v:.1f}×" if v < 1.95 or v > 2.05 else "2.0×",
               title=lambda i, c: (c if i in (0, 5) else f"{c} priority levels"), metric="P99, shortest 50%")
    rows = [[c, x_fmt(R["levels"][k]["short"]["p99"]), x_fmt(R["levels"][k]["all"]["p99"])] for k, c in order]
    leg = ('<div class="legend"><span class="k3"><i></i>Fair share</span><span class="k2"><i></i>Strict priority levels, '
           'fair within a level</span><span class="k1"><i></i>SRPT</span></div>')
    return sc.figure(
        code="FIG. 6 — HOW MANY PRIORITY LEVELS? · W4 · 80% LOAD",
        fid="fig-levels",
        plot=fr.svg(),
        legend_html=leg,
        caption=("<b>Figure 6.</b> 99th-percentile slowdown of the shortest half of messages as the number "
                 "of strict priority levels grows, cutoffs evenly spaced in log(size). Two levels remove "
                 "most of the penalty; four are almost indistinguishable from SRPT."),
        data_html=sc.table("P99 slowdown by number of levels", ["Levels", "Shortest 50%", "All messages"], rows),
    )


# ---------------------------------------------------------------- figure 7 (table)
def fig_workloads():
    names = {"w2": "W2 · Google search", "w3": "W3 · Google, all apps", "w4": "W4 · Facebook Hadoop", "w5": "W5 · DCTCP web search"}
    rows = []
    for w in ("w2", "w3", "w4", "w5"):
        d = R["workloads"][w]
        rows.append([names[w], f"{d['mean_size']:,.0f} B"] +
                    [x_fmt(d[k]["short"]["p99"]) for k in ("ps", "prio8", "srpt", "fifo")])
    tbl = sc.table("P99 slowdown of the shortest 50% of messages, 80% load",
                   ["Workload", "Mean size", "Fair share", "8 levels", "SRPT", "FIFO"], rows)
    return (f'<figure class="fig" id="fig-workloads"><div class="media__bar"><span class="lamp"></span><span class="lamp">'
            f'</span><span class="lamp"></span><span>FIG. 7 — ALL WORKLOADS · 80% LOAD</span></div>'
            f'<div class="fig__plot">{tbl}</div><figcaption><b>Figure 7.</b> The same comparison on four of '
            f"Homa's five workloads (W1 is omitted, as in the Linux paper). Simulated on an idealised link; sizes "
            f"truncated at 1 MB. Mean sizes are from the truncated tables (ATC Table 3 gives 385,315 B for W5)."
            f"</figcaption></figure>")


if __name__ == "__main__":
    for name, fn in [("cdf", fig_cdf), ("exchange", fig_exchange), ("gantt", fig_gantt), ("bysize", fig_bysize),
                     ("byload", fig_byload), ("levels", fig_levels), ("workloads", fig_workloads)]:
        sc.inject(PAGE, name, fn())
        print("built", name)