"""Lab 01 — serving economics: PTUs, self-hosted GPUs, KV cache, batching.

Three questions answered 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?

Integer money (micro-USD) and integer bytes. The simulator advances a tick counter,
never a clock.

    pytest test_lab.py -v
    LAB_MODULE=solution pytest test_lab.py -v    # the reference, must be green
"""

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:
    """``kv_heads`` is separate from attention heads on purpose: grouped-query attention
    shrinks the KV cache by the query:KV head ratio, and that ratio is the single biggest
    lever on how many sequences fit in memory."""

    name: str
    params: int
    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:
        # TODO: every numeric field must be > 0, else ValueError
        raise NotImplementedError

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

        The leading 2 is keys AND values. Be able to recompute this on a whiteboard —
        every concurrency question reduces to it.
        """
        # TODO
        raise NotImplementedError

    def kv_bytes(self, sequence_tokens: int) -> int:
        """Negative -> ValueError; 0 -> 0."""
        # TODO
        raise NotImplementedError

    def weight_bytes(self) -> int:
        # TODO
        raise NotImplementedError


@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 simplification: memory, bandwidth and FLOPs
    scale with ``count``, and the all-reduce cost that does NOT scale is ignored.
    """

    name: str
    memory_bytes: int
    bandwidth_bytes_per_s: int
    flops_per_s: float
    count: int = 1

    def __post_init__(self) -> None:
        # TODO: the three figures > 0, and count > 0
        raise NotImplementedError

    @property
    def total_memory_bytes(self) -> int:
        # TODO
        raise NotImplementedError

    @property
    def total_bandwidth_bytes_per_s(self) -> int:
        # TODO
        raise NotImplementedError

    @property
    def total_flops_per_s(self) -> float:
        # TODO
        raise NotImplementedError


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.

    ``(total_memory - weights) * (1 - overhead_fraction)``, floored at 0 when the model
    does not fit. The reserve is not optional — activations, context and fragmentation
    are real, and over-admitting OOMs the batch, killing in-flight requests.

    ``overhead_fraction`` outside [0, 1) -> ValueError.
    """
    # TODO
    raise NotImplementedError


def max_concurrent_sequences(gpu: GPU, model: ModelShape, sequence_tokens: int,
                             *, overhead_fraction: float = 0.10) -> int:
    """``usable_kv_bytes // kv_bytes(sequence_tokens)``. Zero-length -> ValueError."""
    # TODO
    raise NotImplementedError


# ======================================================================================
# 2. Prefill and decode
# ======================================================================================


def prefill_ms(model: ModelShape, gpu: GPU, prompt_tokens: int) -> float:
    """Compute-bound: ``2 * params * tokens`` FLOPs / total FLOPs per second, in ms.

    (The 2 is a multiply and an add per parameter per token. The quadratic attention term
    is ignored — fine below a few thousand tokens.)
    """
    # TODO
    raise NotImplementedError


def decode_ms_per_token(model: ModelShape, gpu: GPU, *, batch_size: int = 1,
                        context_tokens: int = 0) -> float:
    """Memory-bandwidth-bound.

    Bytes read per step = ``weights + batch_size * kv_bytes(context_tokens)``.
    Per-TOKEN time = (bytes / total bandwidth) / batch_size, in ms.

    The weight read amortizes across the batch; the KV read does not. That asymmetry is
    the entire argument for batching.
    """
    # TODO
    raise NotImplementedError


def arithmetic_intensity(model: ModelShape, batch_size: int, context_tokens: int) -> float:
    """``(2 * params * batch) / (weights + batch * kv_bytes(context))`` — FLOPs per byte."""
    # TODO
    raise NotImplementedError


def ridge_point(gpu: GPU) -> float:
    """FLOPs/byte at which a kernel stops being memory-bound."""
    # TODO
    raise NotImplementedError


# ======================================================================================
# 3. Batching
# ======================================================================================


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

    def __post_init__(self) -> None:
        # TODO: arrival_tick >= 0; both token counts > 0
        raise NotImplementedError


@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:
        # TODO: results per tick; 0.0 when no ticks elapsed
        raise NotImplementedError

    def mean_ttft_ticks(self) -> float:
        # TODO: mean of (first_token_tick - admitted_tick); 0.0 when empty
        raise NotImplementedError

    def mean_latency_ticks(self) -> float:
        # TODO: mean of (finished_tick - admitted_tick); 0.0 when empty
        raise NotImplementedError


@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."""

    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:
        # TODO: sum of kv_bytes(slot.context)
        raise NotImplementedError

    def _projected(self, request: ServingRequest) -> int:
        """KV bytes at the sequence's FINAL length (prompt + output).

        Admitting on the prompt alone over-admits: every sequence grows, and the batch
        OOMs mid-flight, killing requests that were already half-served.
        """
        # TODO
        raise NotImplementedError

    def run(self, requests: Sequence[ServingRequest], *, max_ticks: int = 100_000) -> SimulationReport:
        """Each tick, in this order:

        1. move newly-arrived requests into the queue;
        2. RETIRE finished slots (before admitting — reusing freed capacity in the same
           step is the whole point of continuous batching);
        3. admit from the queue while ``len(slots) < max_batch`` and the projected KV
           fits. A request whose projected KV exceeds the WHOLE budget can never fit:
           reject it rather than queueing it forever;
        4. stop if nothing is running, nothing is queued and nothing is left to arrive;
        5. record peaks, then run one decode step: every slot produces a token and grows
           its context by 1. Set ``first_token_tick`` on a slot's first step.

        Results are returned sorted by request_id so the report is diffable.
        """
        # TODO
        raise NotImplementedError


class StaticBatcher:
    """The thing continuous batching replaced: fill a batch, run it to completion, repeat.

    Every sequence waits for the LONGEST in its batch.
    """

    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:
        """Fill a batch from arrivals (respecting max_batch and the KV budget, rejecting
        anything that can never fit), then advance the clock by the batch's LONGEST
        output length and finish every member at that tick."""
        # TODO
        raise NotImplementedError


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


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


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

    input_micros_per_1k: int
    output_micros_per_1k: int

    def cost_micros(self, input_tokens: int, output_tokens: int) -> int:
        """Divide by 1000 LAST."""
        # TODO
        raise NotImplementedError


@dataclass(frozen=True)
class ProvisionedPricing:
    """``tokens_per_unit_month`` is the number to be suspicious of — it is not a
    datasheet figure, it depends on YOUR input:output mix, and it decides the
    break-even."""

    units: int
    micros_per_unit_month: int
    tokens_per_unit_month: int

    def monthly_cost_micros(self) -> int:
        # TODO
        raise NotImplementedError

    def monthly_capacity_tokens(self) -> int:
        # TODO
        raise NotImplementedError


@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.

    blended_per_1k = c_in*(1 - f) + c_out*f, where f is the OUTPUT fraction.
    break_even_tokens = monthly_cost / blended_per_1k * 1000  (int)
    break_even_utilization = break_even_tokens / capacity (inf when capacity is 0)

    ``output_fraction`` outside [0, 1] -> ValueError.
    """
    # TODO
    raise NotImplementedError


@dataclass(frozen=True)
class SelfHostedPricing:
    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:
        # TODO
        raise NotImplementedError


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.
    Non-positive volume -> ValueError."""
    # TODO
    raise NotImplementedError


@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, then pay per token for the peak.

    The shape almost every mature deployment converges on: size the floor to p50 demand
    so it stays utilized, and let peaks cost money rather than cause 429s — degrade in
    COST rather than in AVAILABILITY.
    """
    # TODO
    raise NotImplementedError


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 volume; return ``(choice, monthly micro-USD)``.

    Ties break toward PAYG — no commitment, no operational burden. (Hint: make the tie
    break explicit by sorting on ``(cost, rank, choice)`` with PAYG at rank 0.)
    """
    # TODO
    raise NotImplementedError


def main() -> None:
    print("implement the TODOs, then compare with `python solution.py`")


if __name__ == "__main__":
    main()