"""Draw the four figures of lesson 30 from rent_vs_own.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_rent_vs_own.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. Rent and owning differ by colour and by line style (a plain line against steps), and the bars say "cheaper" or "dearer" in words.
"""
import math
import sys
from pathlib import Path
from xml.sax.saxutils import escape

import rent_vs_own as rvo

FIGS = Path(__file__).resolve().parent.parent / "figs"
INK, GREY, ORANGE, GREEN, FILL, PANEL = "#0a1626", "#44566b", "#b36200", "#00855F", "#d2dae4", "#ebf0f6"
DASH = "10 8"
WIDTH = 540
TOKENS = 300 * 10**9                                 # the case the trace and the widget open on


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_rent_vs_own.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, halo: bool = False) -> str:
    extra = (f' text-anchor="{anchor}"' if anchor != "start" else "") + (' font-weight="bold"' if bold else "") + (f' font-size="{size}"' if size != 22 else "")
    extra += ' stroke="#ffffff" stroke-width="6" stroke-linejoin="round" paint-order="stroke"' if halo else ""      # white outline, so a line behind the text does not cut through it
    return f'<text x="{num(x)}" y="{num(y)}"{extra} fill="{color}">{escape(content)}</text>'


def money(x) -> str:
    return "$" + rvo.commas(rvo.dollars(x))


def billions(tokens: float) -> str:
    return f"{math.floor(tokens / 1e9 + 0.5):,}"


# --- figure 1: both bills over the month's tokens --------------------------------------------------------------------------------------------

def fig_lines() -> str:
    per_reply, gpu_month, carry = rvo.rates(24)
    step = carry * rvo.TOKENS_PER_REPLY                                            # tokens one planned GPU carries in a month
    cross = rvo.break_even(24)
    now = rvo.case(TOKENS, 24)
    span = 1.25 * TOKENS
    steps = math.ceil(span / step)
    own_at = lambda n: rvo.PEOPLE + (n + rvo.SPARE) * gpu_month
    rent_at = lambda t: per_reply * t / rvo.TOKENS_PER_REPLY
    ymax = math.ceil(max(rent_at(span), own_at(steps)) / 10_000) * 10_000
    x0, x1, top, base = 96, 520, 60, 300
    px = lambda t: x0 + t / span * (x1 - x0)
    py = lambda cost: base - cost / ymax * (base - top)
    body = [text(20, 36, "dollars a month", color=GREY)]
    for tick in range(0, ymax + 1, 20_000):
        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, rvo.commas(tick), "end", GREY, size=20))
    stairs = []
    for k in range(steps):
        stairs += [(px(k * step), py(own_at(k + 1))), (px(min((k + 1) * step, span)), py(own_at(k + 1)))]
    body.append(f'<path d="M{" L".join(f"{num(a)} {num(b)}" for a, b in stairs)}" fill="none" stroke="{INK}" stroke-width="3"/>')
    body.append(f'<path d="M{num(px(0))} {num(py(0))} L{num(px(span))} {num(py(rent_at(span)))}" fill="none" stroke="{ORANGE}" stroke-width="4"/>')
    cx = px(cross["tokens"])
    body.append(f'<path d="M{num(cx)} {base}V{top}" stroke="{GREY}" stroke-width="3" stroke-dasharray="{DASH}"/>')
    body.append(f'<circle cx="{num(cx)}" cy="{num(py(cross["own"]))}" r="8" fill="{GREEN}" stroke="{INK}" stroke-width="2.5"/>')
    body.append(text(cx, top - 10, f"{billions(cross['tokens'])} billion", "middle", INK, bold=True))
    body.append(text(cx - 12, base - 14, "rent costs less", "end", INK, bold=True))
    body.append(text(cx + 12, base - 14, "own costs less", "start", INK, bold=True))
    body.append(text(px(0.85 * span), py(rent_at(0.85 * span)) - 18, "rent", "end", INK, bold=True))
    body.append(text(x1, py(own_at(steps)) + 34, "own", "end", INK, bold=True))
    body.append(f'<path d="M{x0} {base}H{x1}" stroke="{GREY}" stroke-width="2"/>')
    for tick in range(0, int(span / 1e9) + 1, 100):
        body.append(text(px(tick * 1e9), base + 30, str(tick), "middle", GREY, size=20))
    body.append(text((x0 + x1) / 2, base + 62, "billion tokens a month", "middle", GREY))
    desc = (f"Renting (a rising orange line) and owning (black steps) over 0 to {billions(span)} billion tokens a month, in dollars a month, with the traffic spread over 24 hours a day. "
            f"Renting starts at nothing and rises {money(per_reply * 1e9 / rvo.TOKENS_PER_REPLY)} for each billion tokens. Owning starts at {money(own_at(1))} (the people, "
            f"{rvo.cp.gpu_word(1 + rvo.SPARE)}) and steps up {money(gpu_month)} for each GPU added, every {billions(step)} billion tokens or so. The lines cross at "
            f"{billions(cross['tokens'])} billion tokens a month, with {cross['gpus']} GPUs paid for there. At {billions(TOKENS)} billion, renting costs {money(now['rent'])} "
            f"and owning {money(now['own'])}.")
    return frame(WIDTH, 380, f"Rent or own: the lines cross at {billions(cross['tokens'])} billion tokens a month", desc, body)


# --- figure 2: what a fixed set of GPUs costs per million tokens -----------------------------------------------------------------------------

def fig_busy() -> str:
    now = rvo.case(TOKENS, 24)
    fleet = float(now["own"])
    capacity = now["gpus"] * rvo.REPLIES_PER_GPU_SECOND * rvo.SECONDS_PER_MONTH * rvo.TOKENS_PER_REPLY / 1e6     # millions of tokens a month at 100% busy
    rent = float(now["rent_per_million"])
    own_at = lambda busy: fleet / (busy * capacity)
    equal = fleet / (rent * capacity)                                              # the busy share at which this fleet costs what renting costs
    ceiling = rvo.busy_share((now["gpus"] - rvo.SPARE) * rvo.rates(24)[2], now["gpus"])    # the most these GPUs reach: the planned ones at their planned load, the spare idle
    x0, x1, top, base, lo, ymax = 96, 496, 60, 300, 0.10, 0.60
    px = lambda busy: x0 + (busy - lo) / (1 - lo) * (x1 - x0)
    py = lambda cost: base - cost / ymax * (base - top)
    first = math.ceil(fleet / (ymax * capacity) * 100) / 100
    body = [text(20, 36, "dollars per million tokens", color=GREY)]
    for tick in (0, 0.2, 0.4, 0.6):
        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, f"{tick:.2f}", "end", GREY, size=20))
    busy_points = [first + k / 100 for k in range(0, math.floor((ceiling - first) * 100) + 1)] + [ceiling]
    curve = "M" + " L".join(f"{num(px(b))} {num(py(own_at(b)))}" for b in busy_points)
    body.append(f'<path d="{curve}" fill="none" stroke="{INK}" stroke-width="3"/>')
    body.append(f'<path d="M{x0} {num(py(rent))}H{x1}" stroke="{ORANGE}" stroke-width="4"/>')
    body.append(text(x1, py(rent) - 14, f"rent ${rent:.4f}", "end", INK, bold=True))
    body.append(text(px(first) + 14, py(own_at(first)) + 30, f"own, {now['gpus']} GPUs", "start", INK, bold=True))
    body.append(f'<path d="M{num(px(equal))} {base}V{top}" stroke="{GREY}" stroke-width="3" stroke-dasharray="{DASH}"/>')
    body.append(text(px(equal), top - 10, f"{rvo.percent(equal)} busy", "middle", INK, bold=True))
    body.append(f'<circle cx="{num(px(now["busy"]))}" cy="{num(py(own_at(now["busy"])))}" r="8" fill="{GREEN}" stroke="{INK}" stroke-width="2.5"/>')
    body.append(text(px(now["busy"]) - 14, py(own_at(now["busy"])) + 32, f"${own_at(now['busy']):.4f}", "end", INK, halo=True))
    body.append(f'<circle cx="{num(px(ceiling))}" cy="{num(py(own_at(ceiling)))}" r="6" fill="#ffffff" stroke="{INK}" stroke-width="3"/>')
    body.append(text(px(ceiling) + 14, py(own_at(ceiling)) + 32, f"at most {rvo.percent(ceiling)}", "start", INK))
    body.append(f'<path d="M{x0} {base}H{x1}" stroke="{GREY}" stroke-width="2"/>')
    for tick in (0.2, 0.4, 0.6, 0.8, 1.0):
        body.append(text(px(tick), base + 30, f"{round(tick * 100)}%", "middle", GREY, size=20))
    body.append(text((x0 + x1) / 2, base + 62, "how busy the GPUs are", "middle", GREY))
    desc = (f"Dollars per million tokens against how busy the GPUs are, from {round(lo * 100)}% to 100%, for a fixed set of {now['gpus']} GPUs that costs {money(fleet)} a month. "
            f"The black curve is the owning cost per million tokens; it falls as the busy share rises. The orange line is the rent price, ${rent:.4f}. The curve passes under the line at "
            f"{rvo.percent(equal)} busy. At {rvo.percent(now['busy'])} busy, the case of the lesson, owning is ${own_at(now['busy']):.4f}. The curve ends at {rvo.percent(ceiling)} busy, "
            f"the most that {now['gpus'] - rvo.SPARE} planned GPUs and a spare can reach.")
    return frame(WIDTH, 380, f"Owning costs less per million tokens above {rvo.percent(equal)} busy", desc, body)


# --- figure 3: the same tokens, in fewer hours a day -----------------------------------------------------------------------------------------

HOURS_SHOWN = (24, 18, 12, 8, 4)


def fig_peaky() -> str:
    rows = [(h, rvo.case(TOKENS, h)) for h in HOURS_SHOWN]
    rent = float(rows[0][1]["rent"])
    longest = max(float(c["own"]) for _, c in rows)
    x0, span = 24, 250
    scale = span / longest
    body, y = [text(x0, 30, f"dashed mark: renting, {money(rent)} a month", color=GREY)], 76
    for h, c in rows:
        body.append(text(x0, y, f"{h} hours a day: {c['gpus']} GPUs, {rvo.percent(c['busy'])} busy", bold=True))
        own = float(c["own"])
        cheaper = own < rent
        body.append(f'<rect x="{x0}" y="{y + 12}" width="{num(own * scale)}" height="30" fill="{PANEL}" stroke="{GREEN if cheaper else ORANGE}" stroke-width="4"/>')
        body.append(f'<path d="M{num(x0 + rent * scale)} {y + 6}V{y + 48}" stroke="{INK}" stroke-width="3" stroke-dasharray="6 5"/>')
        body.append(text(x0 + max(own, rent) * scale + 10, y + 36, f"{money(own)}, " + ("cheaper" if cheaper else "dearer")))
        y += 68
    facts = "; ".join(f"{h} hours: {c['gpus']} GPUs, {money(c['own'])}, busy {rvo.percent(c['busy'])}" for h, c in rows)
    desc = (f"The same {billions(TOKENS)} billion tokens a month with all the traffic in {', '.join(str(h) for h in HOURS_SHOWN)} hours a day. One bar for what owning costs a month "
            f"against a dashed line at the rent, {money(rent)}. {facts}. Owning is cheaper only when the traffic fills most of the day.")
    return frame(WIDTH, y - 12, "The same tokens cost more to own the fewer hours they arrive in", desc, body)


# --- figure 4: what the people cost does to the crossing -------------------------------------------------------------------------------------

PEOPLE_SHOWN = (0, 5_000, 10_000, 20_000, 30_000)


def fig_people() -> str:
    rows = [(p, rvo.break_even(24, people=p)) for p in PEOPLE_SHOWN]
    longest = max(x["tokens"] for _, x in rows)
    x0, span = 24, 300
    scale = span / longest
    body, y = [], 30
    for people, x in rows:
        body.append(text(x0, y, f"people {money(people)} a month: {x['gpus']} GPUs", bold=True))
        body.append(f'<rect x="{x0}" y="{y + 12}" width="{num(x["tokens"] * scale)}" height="30" fill="{PANEL}" stroke="{INK}" stroke-width="2"/>')
        body.append(text(x0 + x["tokens"] * scale + 10, y + 36, f"{rvo.tokens_words(x['tokens']).replace(' billion', '')} billion"))
        y += 68
    facts = "; ".join(f"{money(p)}: {rvo.tokens_words(x['tokens'])} tokens with {x['gpus']} GPUs" for p, x in rows)
    desc = (f"Where owning first costs no more than renting, by what the people cost a month, with the traffic spread over 24 hours a day. One bar each, longer when the people cost more. {facts}.")
    return frame(WIDTH, y - 10, "The more the people cost, the later owning catches up with renting", desc, body)


FIGURES = {
    "rent-or-own.svg": fig_lines,
    "cost-per-million.svg": fig_busy,
    "same-tokens-fewer-hours.svg": fig_peaky,
    "people-cost.svg": fig_people,
}


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


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