"""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)
               — or the OBSERVED `file_size_gb` of an artifact when one is available (flagged `weights_source: observed`)
    kv_cache = per layer: 2 (K and V) × kv_heads × head_dim × bytes × context × batch — when the architecture metadata is known
               (`num_hidden_layers`, `num_key_value_heads` | `num_attention_heads`, `head_dim` | `hidden_size`); otherwise the documented
               heuristic 0.5 GB per 8 192 tokens (× batch), architecture-agnostic
    fits     = estimated_memory_gb ≤ hardware memory − 2 GB (OS / framework headroom); multi-GPU = sum of device memory, interconnect ignored
"""
from __future__ import annotations

from typing import Any

BYTES_PER_PARAM = {"4bit": 0.5, "8bit": 1.0, "fp16": 2.0, "bf16": 2.0, "fp8": 1.0, "int4": 0.5, "int8": 1.0, "fp32": 4.0}
KV_BYTES = {"fp16": 2.0, "bf16": 2.0, "fp8": 1.0, "int8": 1.0, "8bit": 1.0, "4bit": 2.0, "int4": 2.0, "fp32": 4.0}  # KV cache dtype (4-bit weights usually keep fp16 KV)
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 (or the observed artifact file size when one is recorded).",
    "bytes/param: 4bit = 0.5, 8bit = 1.0, fp16 = 2.0 (uniform quantization, no per-layer exceptions).",
    "KV cache: 2 × layers × kv_heads × head_dim × 2 bytes × context × batch when the architecture is known; otherwise 0.5 GB per 8 192 tokens "
    "(× batch), 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).",
    "Multi-GPU: device memories are summed; interconnect bandwidth, tensor-parallel replication and pipeline bubbles are not modelled.",
]
QUANT_ALIASES = {"q4": "4bit", "int4": "4bit", "nf4": "4bit", "gguf-q4": "4bit", "mlx-4bit": "4bit", "q8": "8bit", "int8": "8bit", "fp8": "8bit",
                 "f16": "fp16", "half": "fp16", "bf16": "fp16", "fp16": "fp16", "4bit": "4bit", "8bit": "8bit", "fp32": "fp32", "f32": "fp32"}


def normalize_quant(q: str | None) -> str:
    s = (q or "4bit").strip().lower()
    return QUANT_ALIASES.get(s, s if s in BYTES_PER_PARAM else "4bit")


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 hardware_memory_options(attrs: dict[str, Any] | None) -> list[float]:
    v = (attrs or {}).get("memory_gb")
    if isinstance(v, list):
        return sorted(float(x) for x in v if isinstance(x, (int, float)) and not isinstance(x, bool))
    m = hardware_memory_gb(attrs)
    return [m] if m is not None else []


def _num(v: Any) -> float | None:
    if isinstance(v, bool):
        return None
    if isinstance(v, (int, float)):
        return float(v)
    if isinstance(v, str):
        try:
            return float(v.replace(",", ""))
        except ValueError:
            return None
    return None


def parameter_count(attrs: dict[str, Any] | None) -> float | None:
    v = _num((attrs or {}).get("parameter_count"))
    return v if v is not None and v > 0 else None


def file_size_gb(attrs: dict[str, Any] | None) -> float | None:
    """Observed weights size (Hugging Face repo file total) when recorded."""
    v = _num((attrs or {}).get("file_size_gb"))
    return v if v is not None and v > 0 else None


def architecture(attrs: dict[str, Any] | None) -> dict[str, float] | None:
    """{layers, kv_heads, head_dim} from Hugging Face-style config keys when all three can be derived, else None."""
    a = attrs or {}
    layers = _num(a.get("num_hidden_layers") or a.get("n_layer") or a.get("num_layers"))
    heads = _num(a.get("num_attention_heads") or a.get("n_head"))
    kv_heads = _num(a.get("num_key_value_heads")) or heads
    head_dim = _num(a.get("head_dim"))
    hidden = _num(a.get("hidden_size") or a.get("n_embd"))
    if head_dim is None and hidden and heads:
        head_dim = hidden / heads
    if layers and kv_heads and head_dim:
        return {"layers": layers, "kv_heads": kv_heads, "head_dim": head_dim}
    return None


def kv_cache_gb(context: int, *, batch: int = 1, quant: str = "4bit", arch: dict[str, float] | None = None) -> tuple[float, str]:
    """(GB, method). `method` is 'architecture' when computed from layers × kv_heads × head_dim, else 'heuristic'."""
    ctx = max(0.0, float(context))
    b = max(1, int(batch))
    if arch:
        bytes_ = KV_BYTES.get(quant, 2.0)
        gb = 2.0 * arch["layers"] * arch["kv_heads"] * arch["head_dim"] * bytes_ * ctx * b / 1e9
        return round(gb, 3), "architecture"
    return round(KV_GB_PER_8K * ctx / 8192.0 * b, 3), "heuristic"


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, "estimated": True,
            "note": f"estimated: {parameter_count_ / 1e9:.1f}B params × {BYTES_PER_PARAM.get(quant, 0.5)} B × 1.15 + KV cache for {context} tokens"}


def fit_detailed(attrs: dict[str, Any] | None, memory_gb: float, *, quant: str = "4bit", context: int = 8192, batch: int = 1, gpu_count: int = 1,
                 observed_size_gb: float | None = None) -> dict[str, Any] | None:
    """Full breakdown (weights / KV cache / overhead / headroom). `observed_size_gb` (artifact file size) replaces the parameter-based weights
    estimate when given. Returns None when neither parameters nor an observed size are known — nothing is estimated from thin air."""
    quant = normalize_quant(quant)
    params = parameter_count(attrs)
    observed = observed_size_gb if observed_size_gb else None
    if params is None and observed is None:
        return None
    if observed is not None:
        weights_raw = observed
        weights_source = "observed"
    else:
        assert params is not None
        weights_raw = params * BYTES_PER_PARAM[quant] / 1e9
        weights_source = "estimated"
    overhead = round(weights_raw * (OVERHEAD - 1), 2)
    kv, kv_method = kv_cache_gb(context, batch=batch, quant=quant, arch=architecture(attrs))
    total = round(weights_raw + overhead + kv, 2)
    total_memory = float(memory_gb) * max(1, int(gpu_count))
    headroom = round(total_memory - RESERVED_GB - total, 2)
    out: dict[str, Any] = {
        "quantization": quant, "estimated": True, "fits": headroom >= 0, "estimated_memory_gb": total, "headroom_gb": headroom,
        "breakdown": {"weights_gb": round(weights_raw, 2), "weights_source": weights_source, "overhead_gb": overhead, "kv_cache_gb": kv, "kv_cache_method": kv_method,
                      "reserved_gb": RESERVED_GB, "context": int(context), "batch": int(batch)},
        "device": {"memory_gb": float(memory_gb), "gpu_count": int(gpu_count), "total_memory_gb": total_memory},
        "note": (f"{'observed' if weights_source == 'observed' else 'estimated'} weights {weights_raw:.1f} GB + 15% overhead + KV cache {kv:.2f} GB ({kv_method}) "
                 f"for {context} tokens × batch {batch}; {RESERVED_GB:.0f} GB reserved for the OS/framework"),
    }
    if params is not None:
        out["parameter_count"] = params
    if gpu_count > 1:
        out["multi_gpu_note"] = "Device memories summed; interconnect bandwidth, tensor-parallel replication and pipeline bubbles are not modelled."
    return out


__all__ = ["ASSUMPTIONS", "BYTES_PER_PARAM", "KV_BYTES", "RESERVED_GB", "architecture", "estimate_memory_gb", "file_size_gb", "fit", "fit_detailed",
           "hardware_memory_gb", "hardware_memory_options", "kv_cache_gb", "normalize_quant", "parameter_count"]