"""Offline illustrative context costs; no API access. Python >=3.11."""
from dataclasses import dataclass
from math import isclose, isfinite
import argparse


def nonnegative(value, name):
    if not isfinite(value) or value < 0:
        raise ValueError(f"{name} must be finite and nonnegative")


@dataclass(frozen=True)
class Scenario:
    old_tokens: int = 120_000
    new_tokens: int = 20_000
    calls: int = 8
    fresh: float = 2.0
    read: float = 0.1
    write: float = 2.5
    compact_cost: float = 0.25

    def __post_init__(self):
        for name in ("old_tokens", "new_tokens", "calls"):
            value = getattr(self, name)
            if isinstance(value, bool) or not isinstance(value, int) or value < 0:
                raise ValueError(f"{name} must be a nonnegative integer")
        if self.calls < 1:
            raise ValueError("calls must be at least one")
        for name in ("fresh", "read", "write", "compact_cost"):
            nonnegative(getattr(self, name), name)

    def kept(self, hit_fraction):
        nonnegative(hit_fraction, "hit_fraction")
        if hit_fraction > 1:
            raise ValueError("hit_fraction must be <=1")
        return self.calls * self.old_tokens / 1_000_000 * (
            hit_fraction * self.read + (1 - hit_fraction) * self.fresh
        )

    def compacted(self):
        return self.compact_cost + self.new_tokens / 1_000_000 * (
            self.write + (self.calls - 1) * self.read
        )


def kv_gib(layers=32, heads=8, head_dim=128, tokens=100_000,
           element_bytes=2, sequences=1):
    values = (layers, heads, head_dim, tokens, element_bytes, sequences)
    if any(isinstance(v, bool) or not isinstance(v, int) or v < 0 for v in values):
        raise ValueError("KV dimensions must be nonnegative integers")
    return 2 * layers * heads * head_dim * tokens * element_bytes * sequences / 2**30


def self_test():
    s = Scenario()
    assert isclose(s.kept(1), 0.096)
    assert isclose(s.kept(0), 1.92)
    assert isclose(s.kept(.9), .2784)
    assert isclose(s.kept(.8), .4608)
    assert isclose(s.compacted(), .314)
    assert isclose(kv_gib(), 12.20703125)
    assert isclose(kv_gib(sequences=4), 48.828125)
    assert Scenario(calls=1, compact_cost=0).compacted() == .05
    for fn in (lambda: s.kept(1.1), lambda: s.kept(float("nan")),
               lambda: Scenario(calls=0), lambda: Scenario(old_tokens=-1),
               lambda: kv_gib(tokens=-1), lambda: Scenario(read=-.1)):
        try:
            fn()
        except ValueError:
            pass
        else:
            raise AssertionError("invalid input accepted")
    print("PASS: arithmetic, one-call boundary, KV dimensions, invalid inputs")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--self-test", action="store_true")
    args = parser.parse_args()
    if args.self_test:
        self_test()
    else:
        scenario = Scenario()
        print(f"Illustrative conventional KV cache: {kv_gib():.2f} GiB")
        print(f"Keep, full cache hits: ${scenario.kept(1):.3f}")
        print(f"Compact, write once then read: ${scenario.compacted():.3f}")
        print(f"Keep, no cache hits: ${scenario.kept(0):.3f}")
        print("Hypothetical rates; excludes unchanged input/output and recovery.")