"""Reference solution — serving economics: PTUs, self-hosted GPUs, KV cache, batching.

Three questions this file answers with arithmetic rather than opinion:

  1. When does dedicated capacity (PTUs / provisioned throughput) beat pay-as-you-go?
  2. How many concurrent sequences fit on a GPU, and why is that a MEMORY question?
  3. What does continuous batching actually buy over static batching?

All integer money (micro-USD) and integer bytes. Deterministic: the simulator advances a
tick counter, never a clock. ``python solution.py`` runs the worked example.
"""

from __future__ import annotations

import math
from dataclasses import dataclass, field, replace
from enum import Enum
from typing import Dict, Iterable, List, Mapping, Optional, Sequence, Tuple

GIB = 1024 ** 3
GB = 1000 ** 3

# ======================================================================================
# 1. Model and hardware shapes
# ======================================================================================


@dataclass(frozen=True)
class ModelShape:
    """Everything about a transformer that determines its serving cost.

    ``kv_heads`` is separate from attention heads on purpose: grouped-query attention
    (GQA) shrinks the KV cache by the ratio of query heads to KV heads, and that ratio is
    the single biggest lever on how many sequences fit in memory.
    """

    name: str
    params: int                 # total parameters
    layers: int
    kv_heads: int
    head_dim: int
    bytes_per_element: int = 2  # fp16/bf16
    bytes_per_param: int = 2

    def __post_init__(self) -> None:
        for field_name in ("params", "layers", "kv_heads", "head_dim",
                           "bytes_per_element", "bytes_per_param"):
            if getattr(self, field_name) <= 0:
                raise ValueError(f"{field_name} must be > 0")

    def kv_bytes_per_token(self) -> int:
        """``2 * layers * kv_heads * head_dim * bytes_per_element``.

        The leading 2 is keys AND values. This is the number to be able to recompute on a
        whiteboard, because every concurrency question reduces to it.
        """
        return 2 * self.layers * self.kv_heads * self.head_dim * self.bytes_per_element

    def kv_bytes(self, sequence_tokens: int) -> int:
        if sequence_tokens < 0:
            raise ValueError("sequence_tokens must be >= 0")
        return self.kv_bytes_per_token() * sequence_tokens

    def weight_bytes(self) -> int:
        return self.params * self.bytes_per_param


@dataclass(frozen=True)
class GPU:
    """One accelerator, or a tensor-parallel group of ``count`` of them.

    Modelling a TP group as one big GPU is a deliberate simplification: memory,
    bandwidth and FLOPs all scale with ``count``, and the all-reduce communication cost
    that does NOT scale is ignored. That overhead is real (single-digit to ~20% of decode
    time depending on interconnect) and it is why TP degree is chosen as small as the
    model allows, not as large as the node permits.
    """

    name: str
    memory_bytes: int           # per accelerator
    bandwidth_bytes_per_s: int  # per accelerator
    flops_per_s: float          # per accelerator, dense, at the model's precision
    count: int = 1

    def __post_init__(self) -> None:
        if self.memory_bytes <= 0 or self.bandwidth_bytes_per_s <= 0 or self.flops_per_s <= 0:
            raise ValueError("GPU figures must be > 0")
        if self.count <= 0:
            raise ValueError("count must be > 0")

    @property
    def total_memory_bytes(self) -> int:
        return self.memory_bytes * self.count

    @property
    def total_bandwidth_bytes_per_s(self) -> int:
        return self.bandwidth_bytes_per_s * self.count

    @property
    def total_flops_per_s(self) -> float:
        return self.flops_per_s * self.count


def usable_kv_bytes(gpu: GPU, model: ModelShape, *, overhead_fraction: float = 0.10) -> int:
    """Memory left for the KV cache after weights and a working-space reserve.

    The reserve is not optional: activations, CUDA context, fragmentation and the
    framework itself all take space, and a planner that assumes 100% of the remainder is
    KV will over-admit and OOM under load — which is the worst possible failure mode
    because it kills in-flight requests, not just the new one.
    """
    if not 0.0 <= overhead_fraction < 1.0:
        raise ValueError("overhead_fraction must be in [0, 1)")
    free = gpu.total_memory_bytes - model.weight_bytes()
    if free <= 0:
        return 0
    return int(free * (1.0 - overhead_fraction))


def max_concurrent_sequences(gpu: GPU, model: ModelShape, sequence_tokens: int,
                             *, overhead_fraction: float = 0.10) -> int:
    """The concurrency ceiling — a MEMORY answer, not a compute one."""
    per_sequence = model.kv_bytes(sequence_tokens)
    if per_sequence <= 0:
        raise ValueError("sequence_tokens must be > 0")
    return usable_kv_bytes(gpu, model, overhead_fraction=overhead_fraction) // per_sequence


# ======================================================================================
# 2. Prefill and decode — two different bottlenecks
# ======================================================================================


def prefill_ms(model: ModelShape, gpu: GPU, prompt_tokens: int) -> float:
    """Compute-bound: roughly ``2 * params * tokens`` FLOPs, processed in parallel.

    The factor 2 is one multiply and one add per parameter per token. Attention adds a
    quadratic term that this ignores — fine below a few thousand tokens, increasingly
    wrong above.
    """
    if prompt_tokens < 0:
        raise ValueError("prompt_tokens must be >= 0")
    flops = 2.0 * model.params * prompt_tokens
    return (flops / gpu.total_flops_per_s) * 1000.0


def decode_ms_per_token(model: ModelShape, gpu: GPU, *, batch_size: int = 1,
                        context_tokens: int = 0) -> float:
    """Memory-bandwidth-bound: each step re-reads the weights ONCE for the whole batch,
    plus each sequence's KV cache.

    This asymmetry is the entire argument for batching: the weight read is amortized
    across the batch, so per-request cost falls as the batch grows — until the KV term
    (which does NOT amortize) dominates.
    """
    if batch_size <= 0:
        raise ValueError("batch_size must be > 0")
    if context_tokens < 0:
        raise ValueError("context_tokens must be >= 0")
    bytes_read = model.weight_bytes() + batch_size * model.kv_bytes(context_tokens)
    seconds = bytes_read / gpu.total_bandwidth_bytes_per_s
    return (seconds / batch_size) * 1000.0


def arithmetic_intensity(model: ModelShape, batch_size: int, context_tokens: int) -> float:
    """FLOPs per byte for one decode step. Compare against the GPU's
    ``flops_per_s / bandwidth_bytes_per_s`` ridge point: below it you are memory-bound.

    Decode at batch 1 is *always* far below the ridge on modern hardware, which is why a
    single-stream self-hosted deployment wastes almost all of its FLOPs.
    """
    if batch_size <= 0:
        raise ValueError("batch_size must be > 0")
    flops = 2.0 * model.params * batch_size
    byts = model.weight_bytes() + batch_size * model.kv_bytes(context_tokens)
    return flops / byts


def ridge_point(gpu: GPU) -> float:
    return gpu.total_flops_per_s / gpu.total_bandwidth_bytes_per_s


# ======================================================================================
# 3. Continuous batching
# ======================================================================================


@dataclass(frozen=True)
class ServingRequest:
    request_id: str
    arrival_tick: int
    prompt_tokens: int
    output_tokens: int

    def __post_init__(self) -> None:
        if self.arrival_tick < 0:
            raise ValueError("arrival_tick must be >= 0")
        if self.prompt_tokens <= 0 or self.output_tokens <= 0:
            raise ValueError("token counts must be > 0")


@dataclass
class ServingResult:
    request_id: str
    admitted_tick: int
    first_token_tick: int
    finished_tick: int
    max_context: int

    @property
    def queue_ticks(self) -> int:
        return self.first_token_tick - self.admitted_tick


@dataclass
class SimulationReport:
    results: List[ServingResult]
    total_ticks: int
    admitted: int
    rejected: List[str]
    peak_batch: int
    peak_kv_bytes: int

    def throughput_per_tick(self) -> float:
        return len(self.results) / self.total_ticks if self.total_ticks else 0.0

    def mean_ttft_ticks(self) -> float:
        if not self.results:
            return 0.0
        return sum(r.first_token_tick - r.admitted_tick for r in self.results) / len(self.results)

    def mean_latency_ticks(self) -> float:
        if not self.results:
            return 0.0
        return sum(r.finished_tick - r.admitted_tick for r in self.results) / len(self.results)


@dataclass
class _Slot:
    request: ServingRequest
    context: int
    produced: int
    admitted_tick: int
    first_token_tick: Optional[int] = None


class ContinuousBatcher:
    """In-flight batching with KV-budget admission control.

    Every tick: retire finished sequences, admit whatever now fits, then run one decode
    step for the whole batch. That is the whole idea — a request does not wait for the
    batch it arrived with to finish, which is why the win over static batching is largest
    exactly when output lengths are uneven.
    """

    def __init__(self, model: ModelShape, gpu: GPU, *, max_batch: int = 256,
                 overhead_fraction: float = 0.10) -> None:
        if max_batch <= 0:
            raise ValueError("max_batch must be > 0")
        self.model = model
        self.gpu = gpu
        self.max_batch = max_batch
        self.kv_budget = usable_kv_bytes(gpu, model, overhead_fraction=overhead_fraction)

    def _kv_used(self, slots: Sequence[_Slot]) -> int:
        return sum(self.model.kv_bytes(s.context) for s in slots)

    def _projected(self, request: ServingRequest) -> int:
        """Admission must budget for the sequence's FINAL length, not its current one.

        Admitting on the prompt length alone over-admits: every sequence grows, and the
        batch OOMs mid-flight, killing requests that were already half-served.
        """
        return self.model.kv_bytes(request.prompt_tokens + request.output_tokens)

    def run(self, requests: Sequence[ServingRequest], *, max_ticks: int = 100_000) -> SimulationReport:
        pending = sorted(requests, key=lambda r: (r.arrival_tick, r.request_id))
        queue: List[ServingRequest] = []
        slots: List[_Slot] = []
        results: List[ServingResult] = []
        rejected: List[str] = []
        peak_batch = 0
        peak_kv = 0
        index = 0
        tick = 0

        while tick < max_ticks:
            while index < len(pending) and pending[index].arrival_tick <= tick:
                queue.append(pending[index])
                index += 1

            # Retire finished sequences BEFORE admitting: the point of continuous
            # batching is that freed capacity is reused in the same step.
            still_running: List[_Slot] = []
            for slot in slots:
                if slot.produced >= slot.request.output_tokens:
                    results.append(ServingResult(
                        request_id=slot.request.request_id,
                        admitted_tick=slot.admitted_tick,
                        first_token_tick=slot.first_token_tick if slot.first_token_tick is not None else tick,
                        finished_tick=tick,
                        max_context=slot.context))
                else:
                    still_running.append(slot)
            slots = still_running

            used = self._kv_used(slots)
            remaining_queue: List[ServingRequest] = []
            for request in queue:
                projected = self._projected(request)
                if projected > self.kv_budget:
                    rejected.append(request.request_id)      # never fits, at any batch size
                    continue
                if len(slots) < self.max_batch and used + projected <= self.kv_budget:
                    slots.append(_Slot(request=request, context=request.prompt_tokens,
                                       produced=0, admitted_tick=tick))
                    used += projected
                else:
                    remaining_queue.append(request)
            queue = remaining_queue

            if not slots and index >= len(pending) and not queue:
                break

            peak_batch = max(peak_batch, len(slots))
            peak_kv = max(peak_kv, self._kv_used(slots))

            for slot in slots:
                if slot.first_token_tick is None:
                    slot.first_token_tick = tick
                slot.produced += 1
                slot.context += 1

            tick += 1

        return SimulationReport(results=sorted(results, key=lambda r: r.request_id),
                                total_ticks=tick, admitted=len(results),
                                rejected=sorted(rejected), peak_batch=peak_batch,
                                peak_kv_bytes=peak_kv)


class StaticBatcher:
    """The thing continuous batching replaced.

    Fill a batch, run it to completion, then start the next. Every sequence in the batch
    waits for the LONGEST one — so a batch containing one 500-token generation and
    fifteen 20-token generations wastes fifteen slots for 480 steps.
    """

    def __init__(self, model: ModelShape, gpu: GPU, *, max_batch: int = 256,
                 overhead_fraction: float = 0.10) -> None:
        self.model = model
        self.gpu = gpu
        self.max_batch = max_batch
        self.kv_budget = usable_kv_bytes(gpu, model, overhead_fraction=overhead_fraction)

    def run(self, requests: Sequence[ServingRequest], *, max_ticks: int = 100_000) -> SimulationReport:
        pending = sorted(requests, key=lambda r: (r.arrival_tick, r.request_id))
        results: List[ServingResult] = []
        rejected: List[str] = []
        peak_batch = 0
        peak_kv = 0
        index = 0
        tick = 0

        while index < len(pending) and tick < max_ticks:
            batch: List[ServingRequest] = []
            used = 0
            while index < len(pending) and len(batch) < self.max_batch:
                candidate = pending[index]
                if candidate.arrival_tick > tick and batch:
                    break
                projected = self.model.kv_bytes(candidate.prompt_tokens + candidate.output_tokens)
                if projected > self.kv_budget:
                    rejected.append(candidate.request_id)
                    index += 1
                    continue
                if used + projected > self.kv_budget:
                    break
                tick = max(tick, candidate.arrival_tick)
                batch.append(candidate)
                used += projected
                index += 1
            if not batch:
                tick += 1
                continue

            peak_batch = max(peak_batch, len(batch))
            peak_kv = max(peak_kv, used)
            start = tick
            longest = max(r.output_tokens for r in batch)
            for request in batch:
                results.append(ServingResult(
                    request_id=request.request_id, admitted_tick=start,
                    first_token_tick=start,
                    finished_tick=start + longest,      # everyone waits for the longest
                    max_context=request.prompt_tokens + request.output_tokens))
            tick = start + longest

        return SimulationReport(results=sorted(results, key=lambda r: r.request_id),
                                total_ticks=tick, admitted=len(results),
                                rejected=sorted(rejected), peak_batch=peak_batch,
                                peak_kv_bytes=peak_kv)


# ======================================================================================
# 4. Capacity economics
# ======================================================================================


class Capacity(str, Enum):
    PAYG = "payg"
    PROVISIONED = "provisioned"
    SELF_HOSTED = "self_hosted"


@dataclass(frozen=True)
class PaygPricing:
    """Per 1 000 tokens, in micro-USD."""

    input_micros_per_1k: int
    output_micros_per_1k: int

    def cost_micros(self, input_tokens: int, output_tokens: int) -> int:
        return (input_tokens * self.input_micros_per_1k
                + output_tokens * self.output_micros_per_1k) // 1000


@dataclass(frozen=True)
class ProvisionedPricing:
    """Dedicated capacity: a fixed monthly cost for a measured throughput.

    ``tokens_per_month`` is the number to be suspicious of — it is not a datasheet
    figure, it depends on YOUR input:output mix, and it is the term that decides the
    break-even.
    """

    units: int
    micros_per_unit_month: int
    tokens_per_unit_month: int

    def monthly_cost_micros(self) -> int:
        return self.units * self.micros_per_unit_month

    def monthly_capacity_tokens(self) -> int:
        return self.units * self.tokens_per_unit_month


@dataclass(frozen=True)
class BreakEven:
    break_even_tokens: int
    break_even_utilization: float
    provisioned_monthly_micros: int
    capacity_tokens: int


def break_even(provisioned: ProvisionedPricing, payg: PaygPricing,
               *, output_fraction: float) -> BreakEven:
    """Where dedicated capacity starts to win.

    ``output_fraction`` is the share of tokens that are OUTPUT — it matters because
    output is several times the price of input, so the blended PAYG rate (and therefore
    the break-even) moves a long way with the traffic mix.
    """
    if not 0.0 <= output_fraction <= 1.0:
        raise ValueError("output_fraction must be in [0, 1]")
    blended_per_1k = (payg.input_micros_per_1k * (1.0 - output_fraction)
                      + payg.output_micros_per_1k * output_fraction)
    if blended_per_1k <= 0:
        raise ValueError("blended PAYG price must be > 0")
    monthly = provisioned.monthly_cost_micros()
    tokens = int(monthly / blended_per_1k * 1000)
    capacity = provisioned.monthly_capacity_tokens()
    return BreakEven(break_even_tokens=tokens,
                     break_even_utilization=tokens / capacity if capacity else math.inf,
                     provisioned_monthly_micros=monthly,
                     capacity_tokens=capacity)


@dataclass(frozen=True)
class SelfHostedPricing:
    """Your own GPUs. The trap is that only the first line is obvious."""

    gpus: int
    micros_per_gpu_hour: int
    hours_per_month: int = 730
    engineering_micros_per_month: int = 0     # the line everyone forgets

    def monthly_cost_micros(self) -> int:
        return (self.gpus * self.micros_per_gpu_hour * self.hours_per_month
                + self.engineering_micros_per_month)


def self_hosted_cost_per_1k_tokens(pricing: SelfHostedPricing, tokens_per_month: int) -> float:
    """Unit cost is a function of UTILIZATION. An idle GPU costs the same as a busy one,
    which is the whole risk of self-hosting."""
    if tokens_per_month <= 0:
        raise ValueError("tokens_per_month must be > 0")
    return pricing.monthly_cost_micros() / (tokens_per_month / 1000)


@dataclass(frozen=True)
class SpilloverPlan:
    provisioned_tokens: int
    payg_tokens: int
    provisioned_micros: int
    payg_micros: int

    @property
    def total_micros(self) -> int:
        return self.provisioned_micros + self.payg_micros


def plan_with_spillover(monthly_tokens: int, provisioned: ProvisionedPricing,
                        payg: PaygPricing, *, output_fraction: float) -> SpilloverPlan:
    """Fill the dedicated floor first, spill the remainder to pay-as-you-go.

    This is the shape almost every mature deployment converges on: size the floor to p50
    demand so it stays well utilized, and let the peaks cost money rather than cause
    429s. You degrade in COST rather than in AVAILABILITY.
    """
    if monthly_tokens < 0:
        raise ValueError("monthly_tokens must be >= 0")
    capacity = provisioned.monthly_capacity_tokens()
    on_ptu = min(monthly_tokens, capacity)
    spill = monthly_tokens - on_ptu
    spill_input = int(spill * (1.0 - output_fraction))
    spill_output = spill - spill_input
    return SpilloverPlan(
        provisioned_tokens=on_ptu, payg_tokens=spill,
        provisioned_micros=provisioned.monthly_cost_micros(),
        payg_micros=payg.cost_micros(spill_input, spill_output))


def cheapest_option(monthly_tokens: int, *, provisioned: ProvisionedPricing,
                    payg: PaygPricing, self_hosted: Optional[SelfHostedPricing] = None,
                    output_fraction: float) -> Tuple[Capacity, int]:
    """Compare all three at a given volume. Returns ``(choice, monthly micro-USD)``.

    Ties break toward PAYG — the option with no commitment and no operational burden.
    """
    payg_input = int(monthly_tokens * (1.0 - output_fraction))
    payg_output = monthly_tokens - payg_input
    options: List[Tuple[int, int, Capacity]] = [
        (payg.cost_micros(payg_input, payg_output), 0, Capacity.PAYG),
        (plan_with_spillover(monthly_tokens, provisioned, payg,
                             output_fraction=output_fraction).total_micros, 1,
         Capacity.PROVISIONED),
    ]
    if self_hosted is not None:
        options.append((self_hosted.monthly_cost_micros(), 2, Capacity.SELF_HOSTED))
    cost, _, choice = min(options)
    return choice, cost


# ======================================================================================
# Worked example
# ======================================================================================


def _fmt_bytes(n: int) -> str:
    for unit, size in (("GiB", GIB), ("MiB", 1024 ** 2), ("KiB", 1024)):
        if n >= size:
            return f"{n / size:.2f} {unit}"
    return f"{n} B"


def _usd(micros: int) -> str:
    return f"${micros / 1_000_000:,.2f}"


def main() -> None:  # pragma: no cover - narrative output
    model = ModelShape(name="llama-open-70b-gqa", params=70_000_000_000, layers=80,
                       kv_heads=8, head_dim=128)
    one = GPU(name="80GB accelerator", memory_bytes=80 * GIB,
              bandwidth_bytes_per_s=3_350 * GB, flops_per_s=989e12)
    gpu = replace(one, name="8x80GB node (TP=8)", count=8)

    print("=" * 78)
    print("1. THE KV CACHE IS THE CONCURRENCY LIMIT")
    print("=" * 78)
    print(f"  model            : {model.name}, {model.layers} layers, "
          f"{model.kv_heads} KV heads, head_dim {model.head_dim}")
    print(f"  weights          : {_fmt_bytes(model.weight_bytes())}")
    print(f"  one 80 GiB card  : usable for KV = "
          f"{_fmt_bytes(usable_kv_bytes(one, model))}  -> the model does not even fit;")
    print(f"                     tensor parallelism is not an optimization here, it is a")
    print(f"                     precondition. Serving on {gpu.name}:")
    print(f"  KV per token     : {model.kv_bytes_per_token():,} bytes")
    for seq in (2_048, 8_192, 32_768):
        print(f"  KV @ {seq:>6} tok  : {_fmt_bytes(model.kv_bytes(seq)):>10}  ->  "
              f"{max_concurrent_sequences(gpu, model, seq):>3} concurrent sequences")
    print(f"  usable for KV    : {_fmt_bytes(usable_kv_bytes(gpu, model))} "
          f"(640 GiB - weights - 10% working space)")
    print("  -> max concurrency is a MEMORY question, not a compute one.")

    mqa = replace(model, kv_heads=1, name="same model, MQA")
    mha = replace(model, kv_heads=64, name="same model, MHA")
    print()
    print("  the GQA lever, same model, three attention layouts @ 8k context:")
    for variant in (mha, model, mqa):
        print(f"      kv_heads={variant.kv_heads:>3}  KV/token {variant.kv_bytes_per_token():>7,} B"
              f"  ->  {max_concurrent_sequences(gpu, variant, 8192):>4} sequences")

    print()
    print("=" * 78)
    print("2. PREFILL IS COMPUTE-BOUND, DECODE IS MEMORY-BOUND")
    print("=" * 78)
    print(f"  GPU ridge point  : {ridge_point(gpu):.1f} FLOPs/byte")
    print(f"  prefill 2 000 tok: {prefill_ms(model, gpu, 2000):.1f} ms")
    print("  decode, per output token, as the batch grows:")
    for batch in (1, 8, 32, 128):
        per_token = decode_ms_per_token(model, gpu, batch_size=batch, context_tokens=2000)
        intensity = arithmetic_intensity(model, batch, 2000)
        print(f"      batch {batch:>4}: {per_token:6.2f} ms/token   "
              f"arithmetic intensity {intensity:6.2f} FLOPs/byte"
              f"{'  (memory-bound)' if intensity < ridge_point(gpu) else '  (compute-bound)'}")
    print("  -> the weight read is amortized across the batch; the KV read is not.")
    print("  -> this asymmetry is the entire argument for batching.")

    print()
    print("=" * 78)
    print("3. CONTINUOUS VS STATIC BATCHING")
    print("=" * 78)
    # Uneven output lengths: the case static batching handles worst.
    requests = []
    for i in range(24):
        requests.append(ServingRequest(f"r{i:02d}", arrival_tick=i // 4,
                                       prompt_tokens=512,
                                       output_tokens=400 if i % 8 == 0 else 20))
    cont = ContinuousBatcher(model, gpu, max_batch=8).run(requests)
    static = StaticBatcher(model, gpu, max_batch=8).run(requests)
    print(f"  {'':<22}{'continuous':>12}{'static':>12}")
    print(f"  {'total ticks':<22}{cont.total_ticks:>12}{static.total_ticks:>12}")
    print(f"  {'mean latency (ticks)':<22}{cont.mean_latency_ticks():>12.1f}"
          f"{static.mean_latency_ticks():>12.1f}")
    print(f"  {'mean TTFT (ticks)':<22}{cont.mean_ttft_ticks():>12.1f}"
          f"{static.mean_ttft_ticks():>12.1f}")
    print(f"  {'peak batch':<22}{cont.peak_batch:>12}{static.peak_batch:>12}")
    speedup = static.total_ticks / cont.total_ticks if cont.total_ticks else 0
    print(f"  -> continuous finishes the same work {speedup:.2f}x sooner, because a short")
    print("     generation retires immediately instead of waiting for the longest in its batch.")

    print()
    print("=" * 78)
    print("4. ADMISSION CONTROL")
    print("=" * 78)
    tiny = GPU("single 24 GiB card", memory_bytes=24 * GIB,
               bandwidth_bytes_per_s=1_000 * GB, flops_per_s=100e12)
    small_model = ModelShape("small-7b", params=7_000_000_000, layers=32,
                             kv_heads=8, head_dim=128)
    batcher = ContinuousBatcher(small_model, tiny, max_batch=64)
    print(f"  KV budget        : {_fmt_bytes(batcher.kv_budget)}")
    print(f"  a 200k-token request needs {_fmt_bytes(small_model.kv_bytes(200_000))}")
    report = batcher.run([
        ServingRequest("ok-1", 0, 1_000, 100),
        ServingRequest("too-big", 0, 100_000, 100_000),
        ServingRequest("ok-2", 0, 1_000, 100),
    ])
    print(f"  admitted={report.admitted}  rejected={report.rejected}")
    print("  -> a request that can never fit is rejected, not queued forever.")
    print("  -> admission budgets the FINAL length, or the batch OOMs mid-flight.")

    print()
    print("=" * 78)
    print("5. PTU VS PAY-AS-YOU-GO")
    print("=" * 78)
    payg = PaygPricing(input_micros_per_1k=3_000, output_micros_per_1k=12_000)
    ptu = ProvisionedPricing(units=100, micros_per_unit_month=120_000_000,
                             tokens_per_unit_month=40_000_000)
    for mix in (0.1, 0.25, 0.5):
        be = break_even(ptu, payg, output_fraction=mix)
        print(f"  output fraction {mix:>4.0%}: break-even at {be.break_even_tokens / 1e9:5.2f}B tokens"
              f"  = {be.break_even_utilization:5.1%} of the {be.capacity_tokens / 1e9:.1f}B capacity")
    print(f"  monthly commitment : {_usd(ptu.monthly_cost_micros())}")
    print("  -> the break-even moves a LONG way with the traffic mix, because output")
    print("     costs several times input. Never quote a single number.")

    print()
    print("  spillover: fill the floor, then pay per token for the peak")
    for volume_b in (2, 4, 6):
        volume = int(volume_b * 1e9)
        plan = plan_with_spillover(volume, ptu, payg, output_fraction=0.25)
        pure_payg = payg.cost_micros(int(volume * 0.75), int(volume * 0.25))
        print(f"      {volume_b}B tokens: PTU {_usd(plan.provisioned_micros)} + "
              f"spill {_usd(plan.payg_micros)} = {_usd(plan.total_micros)}"
              f"   (pure PAYG {_usd(pure_payg)})")

    print()
    print("=" * 78)
    print("6. AND SELF-HOSTING")
    print("=" * 78)
    hosted = SelfHostedPricing(gpus=8, micros_per_gpu_hour=3_000_000,
                               engineering_micros_per_month=25_000_000)
    print(f"  8 GPUs @ $3/h + $25k/mo engineering = {_usd(hosted.monthly_cost_micros())}/mo")
    for volume_b in (1, 4, 20):
        volume = int(volume_b * 1e9)
        unit = self_hosted_cost_per_1k_tokens(hosted, volume)
        choice, cost = cheapest_option(volume, provisioned=ptu, payg=payg,
                                       self_hosted=hosted, output_fraction=0.25)
        print(f"      {volume_b:>2}B tokens/mo: self-hosted unit cost "
              f"{unit:7.1f} micro-USD/1k   cheapest = {choice.value} ({_usd(cost)})")
    print("  -> self-hosting is a UTILIZATION bet: an idle GPU costs the same as a busy one.")
    print("  -> and the engineering line is real; leaving it out is how the business case")
    print("     for self-hosting is usually won and then lost.")


if __name__ == "__main__":  # pragma: no cover
    main()
