"""ESTIMATED memory footprint of a model on a piece of hardware. Transparent formula, always labelled as an estimate:

    weights  = parameter_count × bytes_per_param × 1.15   (runtime overhead: activations, buffers, fragmentation)
    kv_cache = 0.5 GB per 8 192 tokens of context          (coarse, architecture-agnostic)
    fits     = estimated_memory_gb ≤ hardware memory − 2 GB (OS / framework headroom)
"""
from __future__ import annotations

from typing import Any

BYTES_PER_PARAM = {"4bit": 0.5, "8bit": 1.0, "fp16": 2.0}
OVERHEAD = 1.15
KV_GB_PER_8K = 0.5
RESERVED_GB = 2.0
ASSUMPTIONS = [
    "Estimated, not measured: weights = parameters × bytes/param × 1.15 runtime overhead.",
    "bytes/param: 4bit = 0.5, 8bit = 1.0, fp16 = 2.0 (uniform quantization, no per-layer exceptions).",
    "KV cache approximated at 0.5 GB per 8 192 tokens of context, independent of architecture (GQA/MLA models need less).",
    "A model 'fits' when the estimate is at most the device memory minus 2 GB reserved for the OS and framework.",
    "Mixture-of-experts models are estimated on total parameters (all experts must be resident); active parameters are ignored.",
    "Device memory uses the largest configuration when several are listed (e.g. Apple silicon tiers).",
]


def estimate_memory_gb(parameter_count: float, quant: str = "4bit", context: int = 8192) -> float:
    bpp = BYTES_PER_PARAM.get(quant, BYTES_PER_PARAM["4bit"])
    weights = parameter_count * bpp * OVERHEAD / 1e9
    kv = KV_GB_PER_8K * max(0.0, float(context)) / 8192.0
    return round(weights + kv, 2)


def hardware_memory_gb(attrs: dict[str, Any] | None) -> float | None:
    """`memory_gb` may be a number or a list of configurations — use the largest."""
    v = (attrs or {}).get("memory_gb")
    if isinstance(v, list):
        nums = [float(x) for x in v if isinstance(x, (int, float)) and not isinstance(x, bool)]
        return max(nums) if nums else None
    if isinstance(v, (int, float)) and not isinstance(v, bool):
        return float(v)
    if isinstance(v, str):
        try:
            return float(v)
        except ValueError:
            return None
    return None


def parameter_count(attrs: dict[str, Any] | None) -> float | None:
    v = (attrs or {}).get("parameter_count")
    if isinstance(v, (int, float)) and not isinstance(v, bool) and v > 0:
        return float(v)
    if isinstance(v, str):
        try:
            f = float(v)
            return f if f > 0 else None
        except ValueError:
            return None
    return None


def fit(parameter_count_: float, memory_gb: float, quant: str = "4bit", context: int = 8192) -> dict[str, Any]:
    est = estimate_memory_gb(parameter_count_, quant, context)
    headroom = round(memory_gb - RESERVED_GB - est, 2)
    return {"quantization": quant, "estimated_memory_gb": est, "fits": headroom >= 0, "headroom_gb": headroom,
            "note": f"estimated: {parameter_count_ / 1e9:.1f}B params × {BYTES_PER_PARAM.get(quant, 0.5)} B × 1.15 + KV cache for {context} tokens"}


__all__ = ["ASSUMPTIONS", "BYTES_PER_PARAM", "estimate_memory_gb", "fit", "hardware_memory_gb", "parameter_count"]