"""Rent a model API or run your own: each way's bill for a month, and the monthly token count where they cross.

Python standard library only. Run `python3 show_numbers.py` to print every number the lesson quotes, and `python3 trace.py` to follow one case step by step.

The model: a service sends a number of tokens a month, all of it in the same hours every day. Renting bills each token at its price (lesson 7's bill). Owning pays
for enough GPUs for the busiest hour (lesson 29's plan) plus a spare, for every hour of the month, plus a fixed monthly cost for the people who run them.
Prices are dollars; token prices are per million tokens.
"""
from __future__ import annotations

import importlib.util
import math
import statistics
import sys
from decimal import ROUND_HALF_UP, Decimal
from pathlib import Path

sys.dont_write_bytecode = True                       # loading the sibling lessons' files below must not leave a __pycache__ in their folders
LESSONS = Path(__file__).resolve().parents[2]


def _load(name: str, folder: str, filename: str):
    """A sibling lesson folder is not a package, so its file is loaded by path."""
    spec = importlib.util.spec_from_file_location(name, LESSONS / folder / "code" / filename)
    module = importlib.util.module_from_spec(spec)
    sys.modules[spec.name] = module                  # dataclasses looks the module up by name
    spec.loader.exec_module(module)
    return module


bill = _load("bill_lesson_7", "why-ai-bill-is-high", "bill.py")
cp = _load("capacity_plan_lesson_29", "how-many-gpus-for-1000-users", "capacity_plan.py")

# --- the inputs: round toy figures, not quotes -----------------------------------------------------------------------------------------------
MIX = cp.OPENING_MIX                                 # lesson 29's usual chats
TARGET_MS = cp.TARGET_MS                             # and its target: the p99 gap between words stays under 50 ms
HEADROOM = cp.HEADROOM                               # and its 20% in hand
DAYS = bill.DAYS_PER_MONTH                           # lesson 7's month: 30 days
HOURS_PER_MONTH = 24 * DAYS
SECONDS_PER_MONTH = HOURS_PER_MONTH * 3600
_chats = cp.chats_of(MIX)
TOKENS_IN = round(statistics.fmean(_chats["prompts"]))      # 1,882: the prompt each reply is sent with, on average
TOKENS_OUT = round(statistics.fmean(_chats["replies"]))     # 157: the reply itself
TOKENS_PER_REPLY = TOKENS_IN + TOKENS_OUT
PRICE_IN = Decimal("0.10")                           # dollars per million input tokens: the cheapest row on the price page
PRICE_OUT = Decimal("0.50")                          # dollars per million output tokens
GPU_HOUR = Decimal("3.00")                           # dollars to keep one GPU for one hour
PEOPLE = 10_000                                      # dollars a month for the people who run the GPUs
SPARE = 1                                            # lesson 29's spare
MAX_GPUS = 5000                                      # the search for a break-even stops here: past this the toy has nothing more to say

DESIGN = cp.plan(1, 100, 3600, MIX, TARGET_MS)                               # lesson 29's plan, for the numbers every GPU count here is divided by
REPLIES_PER_GPU_SECOND = DESIGN["full"] / DESIGN["stay"]  # replies a second one GPU finishes when it is as full as its tighter limit allows

OPENING = {"tokens_b": 300, "hours": 24, "people": PEOPLE, "price_in": 0.10, "price_out": 0.50, "gpu_hour": 3.0, "spare": SPARE}
TOKENS_B = [1, 2, 5, 10, 20, 30, 50, 75, 100, 150, 200, 300, 400, 500, 750, 1000, 2000]    # the widget's monthly tokens, in billions
PEOPLE_RANGE = (0, 40_000, 1000)                     # lowest, highest, step
HOURS_RANGE = (1, 24, 1)
PRICES_IN = [0.10, 0.50, 1.0, 2.0, 4.0, 10.0]
PRICES_OUT = [0.50, 2.0, 5.0, 10.0, 20.0, 50.0]
GPU_HOURS = [1.0, 1.5, 2.0, 2.5, 3.0, 4.0, 5.0, 6.0, 8.0]
SPARES = [0, 1, 2]


# --- the six steps ---------------------------------------------------------------------------------------------------------------------------

def replies_for(tokens):
    return math.floor(tokens / TOKENS_PER_REPLY + 0.5)


def monthly_rent(replies, price_in=PRICE_IN, price_out=PRICE_OUT):
    if replies < 1:
        return Decimal(0)
    return bill.bill(bill.Chat(system=0, message=TOKENS_IN, reply=TOKENS_OUT, turns=replies, keep=0), price_in, price_out).total


def gpus_to_pay_for(replies, hours, spare=SPARE):
    return cp.plan(1, 100, replies / (DAYS * hours), MIX, TARGET_MS)["n"] + spare


def monthly_own(gpus, gpu_hour=GPU_HOUR, people=PEOPLE):
    return gpus * HOURS_PER_MONTH * gpu_hour + people


def busy_share(replies, gpus):
    return replies / (gpus * REPLIES_PER_GPU_SECOND * SECONDS_PER_MONTH)


def rates(hours, price_in=PRICE_IN, price_out=PRICE_OUT, gpu_hour=GPU_HOUR):
    """What the owning and renting sums are made of: dollars the rent charges for one reply, dollars one GPU costs for the month, and the replies one planned GPU
    carries in a month when all the traffic falls in `hours` hours a day (HEADROOM of what it finishes at its limit)."""
    per_reply = (TOKENS_IN * round(price_in * 100) + TOKENS_OUT * round(price_out * 100)) / 10**8
    gpu_month = round(gpu_hour * 100) * HOURS_PER_MONTH / 100
    carry = HEADROOM * REPLIES_PER_GPU_SECOND * DAYS * hours * 3600
    return per_reply, gpu_month, carry


def break_even(hours, price_in=PRICE_IN, price_out=PRICE_OUT, gpu_hour=GPU_HOUR, people=PEOPLE, spare=SPARE):
    """The monthly replies and tokens at which owning first costs no more than renting, with the GPUs paid for there, or None when that takes over MAX_GPUS GPUs.

    With n planned GPUs the owning bill is the same for every reply count up to n * carry, while the rent keeps rising. So the first n whose bill is no more than the
    rent at n * carry holds the crossing, and the crossing is where the rent has risen to that bill.
    """
    per_reply, gpu_month, carry = rates(hours, price_in, price_out, gpu_hour)
    for n in range(1, MAX_GPUS + 1):
        own = people + (n + spare) * gpu_month
        if own <= per_reply * n * carry:
            return {"replies": own / per_reply, "tokens": own / per_reply * TOKENS_PER_REPLY, "gpus": n + spare, "own": own}
    return None


# --- around the steps ------------------------------------------------------------------------------------------------------------------------

def dollars(x) -> int:
    """Whole dollars, halves up."""
    return int(Decimal(x).quantize(Decimal(1), rounding=ROUND_HALF_UP))


def case(tokens, hours, people=PEOPLE, price_in=PRICE_IN, price_out=PRICE_OUT, gpu_hour=GPU_HOUR, spare=SPARE) -> dict:
    """Both bills for one month, the GPUs and the busy share, from the six steps."""
    replies = replies_for(tokens)
    rent = monthly_rent(replies, price_in, price_out)
    gpus = gpus_to_pay_for(replies, hours, spare)
    own = monthly_own(gpus, gpu_hour, people)
    return {"tokens": tokens, "replies": replies, "rent": rent, "gpus": gpus, "own": own, "busy": busy_share(replies, gpus),
            "own_per_million": own / Decimal(tokens) * bill.PER, "rent_per_million": rent / Decimal(tokens) * bill.PER}


def commas(n: int) -> str:
    return f"{n:,}"


def tokens_words(tokens: float) -> str:
    """A token count to three figures, in millions, billions or trillions."""
    unit, name = (1e12, "trillion") if tokens >= 1e12 else (1e9, "billion") if tokens >= 1e9 else (1e6, "million")
    v = tokens / unit
    places = 0 if v >= 100 else 1 if v >= 10 else 2
    return f"{math.floor(v * 10**places + 0.5) / 10**places:.{places}f} {name}"


def percent(x: float) -> str:
    return f"{math.floor(x * 100 + 0.5)}%"


# --- the widget: the words and the chart for one setting of its controls ---------------------------------------------------------------------

W, H, PAD = 300, 150, 4                              # the chart's drawing box; the lines use all of it, and the words under it say what the sides are
MAX_STEPS = 120                                      # the chart draws at most this many steps of the owning line, grouping GPUs when there are more


def fix1(x: float) -> float:
    return math.floor(x * 10 + 0.5) / 10


def spare_words(spare: int) -> str:
    return "no spare" if spare == 0 else f"{spare} spare" if spare == 1 else f"{spare} spares"


def chart(hours, tokens, cross, people, price_in, price_out, gpu_hour, spare, rent_now, own_now) -> dict:
    """The two lines of the chart, as path text for an SVG, over the monthly tokens from 0 to a little past both the reader's setting and the crossing, and the
    heights of the two bills at the reader's setting."""
    per_reply, gpu_month, carry = rates(hours, price_in, price_out, gpu_hour)
    span = 1.25 * max(tokens, cross["tokens"] if cross else 0)
    step = carry * TOKENS_PER_REPLY                        # tokens one planned GPU carries in a month
    steps = math.ceil(span / step)
    stride = max(1, math.ceil(steps / MAX_STEPS))
    own_at = lambda n: people + (n + spare) * gpu_month
    top = max(per_reply * span / TOKENS_PER_REPLY, own_at(steps))
    x = lambda t: fix1(PAD + (W - 2 * PAD) * t / span)
    y = lambda cost: fix1(H - PAD - (H - 2 * PAD) * cost / top)
    points = []
    for k in range(0, steps, stride):
        n = min(k + stride, steps)
        points += [(x(k * step), y(own_at(n))), (x(min(n * step, span)), y(own_at(n)))]
    joined = lambda pts: "M" + " L".join(f"{px:.1f},{py:.1f}" for px, py in pts)
    return {"rent": joined([(x(0), y(0)), (x(span), y(per_reply * span / TOKENS_PER_REPLY))]), "own": joined(points),
            "cross": None if cross is None else f"{x(cross['tokens']):.1f}", "mark": f"{x(tokens):.1f}", "dots": [f"{y(rent_now):.1f}", f"{y(own_now):.1f}"],
            "caption": f"Across: 0 to {tokens_words(span)} tokens a month. Up: 0 to ${commas(dollars(top))} a month."}


def describe(tokens_b, hours, people, price_in, price_out, gpu_hour, spare) -> dict:
    """Everything the widget shows for one setting of its controls: three lines of words and the two lines of the chart."""
    price_in, price_out, gpu_hour = Decimal(str(price_in)), Decimal(str(price_out)), Decimal(str(gpu_hour))
    tokens = tokens_b * 10**9
    now = case(tokens, hours, people, price_in, price_out, gpu_hour, spare)
    cross = break_even(hours, price_in, price_out, gpu_hour, people, spare)
    rent, own = dollars(now["rent"]), dollars(now["own"])
    cheaper = "Owning costs less." if now["own"] < now["rent"] else "Renting costs less." if now["rent"] < now["own"] else "They cost the same."
    return {"headline": f"Below about {tokens_words(cross['tokens'])} tokens a month, renting costs less." if cross else "At these settings owning never catches up with renting.",
            "verdict": f"At {commas(tokens_b)} billion tokens a month, renting costs ${commas(rent)} and owning ${commas(own)}. {cheaper}",
            "detail": f"Owning takes {cp.gpu_word(now['gpus'])} ({spare_words(spare)}), busy {percent(now['busy'])} of the time.",
            "chart": chart(hours, tokens, cross, people, float(price_in), float(price_out), float(gpu_hour), spare, float(now["rent"]), float(now["own"]))}