"""
Speed estimation: how fast will it actually feel?

Two-tier design, with provenance the UI always shows:

  1. TRAINED MODEL (when present): an XGBoost regressor trained on real
     community measurements (LocalScore; 6,633 training rows derived from
     those sources), following the
     methodology of LLM-Pilot (IBM, SC'24, arXiv:2410.02425 — gradient
     boosting over hardware+model features, grouped k-fold by accelerator
     label; NOT a strict leave-one-hardware-out split, ~48% alias overlap).
     Loaded from model/speed_model.skops if scripts/train_speed_model.py
     has been run. method = "measured-model".
  2. ROOFLINE BASELINE (always available, fully offline): decode is memory-
     bandwidth-bound — tok/s ~ bandwidth / bytes-read-per-token (weights +
     KV), times an empirical efficiency factor. See kipply's "Transformer
     Inference Arithmetic" and the JAX scaling book inference chapter.
     method = "roofline".

The anti-gimmick rule lives in the training script: the trained model ships
only if it beats this baseline on held-out hardware; otherwise the baseline
IS the product and the UI says so.

Scope note (honest): this predicts LLM/VLM decode speed. Vision (YOLO) and
diffusion models are COMPUTE-bound, not bandwidth-bound — FPS scales with
TFLOPS / model GFLOPs, a different axis with different data (Ultralytics
publishes per-size GFLOPs and official T4 latencies; dbgpu has per-GPU
TFLOPS). That path is designed in SPEED-BRICK-RESEARCH.md §8 but not built;
non-LLM families keep their provenance-labelled memory verdicts only, rather
than getting fake speed numbers.
"""

import json
import re
from functools import lru_cache
from pathlib import Path

_ROOT = Path(__file__).resolve().parent.parent
_SPECS_PATH = _ROOT / "data" / "gpu_specs.json"
_MODEL_PATH = _ROOT / "model" / "speed_model.skops"

# Decode efficiency vs theoretical bandwidth roofline. Real stacks land well
# under the ceiling; 0.55-0.70 is the typical consumer-GPU range in community
# measurements. We centre conservatively and report a band, never a point.
_EFF_MID, _EFF_LO, _EFF_HI = 0.60, 0.42, 0.78
# Fixed per-token overhead (seconds): kernel launches, sampling, framework cost.
# Without this, the pure bandwidth roofline says a 0.5B model does ~1500 tok/s,
# which is nonsense — single-stream decode is overhead-bound once the model is
# tiny, so tok/s plateaus (~250-350 on fast GPUs) instead of growing forever.
_TOKEN_OVERHEAD_S = 0.0032
# Conservative system-RAM bandwidth for offload modelling (dual-channel DDR4/5).
_RAM_BW_GBS = 48.0
# Reading speed reference: ~4.5 words/s, ~0.75 words per token -> ~6 tok/s.
_READING_TPS = 6.0


@lru_cache(maxsize=1)
def _specs() -> dict:
    try:
        return json.loads(_SPECS_PATH.read_text(encoding="utf-8"))
    except OSError:
        return {"gpus": {}, "apple": {}, "sbc": {}}


def _norm(s: str) -> str:
    return re.sub(r"\s+", " ", re.sub(r"[^a-z0-9 ]", " ", (s or "").lower())).strip()


def _key_matches(key: str, text: str) -> bool:
    """Token-boundary containment: the key's whole tokens must appear as a
    contiguous run in text. Plain substring matching (the old `key in text`)
    let a short key match inside a longer token — e.g. Apple "m2" matched
    "Tesla V100-SXM2" -> "tesla v100 sxm2" (the 'sxm2' token), assigning the
    V100 Apple's bandwidth. Token matching kills that without losing real hits
    like "4080 super" inside "nvidia geforce rtx 4080 super"."""
    kt = key.split()
    tt = text.split()
    n = len(kt)
    if not n:
        return False
    return any(tt[i:i + n] == kt for i in range(len(tt) - n + 1))


@lru_cache(maxsize=1)
def _bw_index() -> tuple:
    idx = []
    for name, d in _specs()["gpus"].items():
        idx.append((_norm(name), float(d["bw"]), float(d.get("vram", 0))))
    idx.sort(key=lambda t: -len(t[0]))   # longest first: '4080 super' beats '4080'
    return tuple(idx)


# Apple chips: the UI only knows base/Pro/Max/Ultra, not the generation. We use
# M2-generation numbers as the conservative representative (older = slower).
_APPLE_TIER_BW = None


def _apple_bw(tier_hint: str) -> float:
    global _APPLE_TIER_BW
    if _APPLE_TIER_BW is None:
        a = {k: v["bw"] for k, v in _specs()["apple"].items()}
        _APPLE_TIER_BW = {
            "ultra": a.get("m2 ultra") or a.get("m1 ultra") or 800.0,
            "max": a.get("m2 max") or 400.0,
            "pro": a.get("m2 pro") or 200.0,
            "base": a.get("m2") or 100.0,
        }
    t = (tier_hint or "").lower()
    for key in ("ultra", "max", "pro"):
        if key in t:
            return _APPLE_TIER_BW[key]
    return _APPLE_TIER_BW["base"]


def bandwidth_for_spec(spec, gpu_label: str = "") -> tuple[float | None, str]:
    """(memory bandwidth GB/s on the fast path, source-note) for a machine."""
    if spec.is_apple_silicon:
        return _apple_bw(gpu_label or spec.gpu_label), "Apple unified memory (conservative M2-gen figure)"
    if spec.gpu_vendor in ("nvidia", "amd", "intel") and spec.vram_gb > 0:
        n = _norm(gpu_label or spec.gpu_label)
        # pass 1: name + VRAM proximity (disambiguates 8 vs 16 GB variants);
        # pass 2: name only — a custom VRAM override must not hide the chart.
        for check_vram in (True, False):
            for key, bw, vram in _bw_index():
                if key and _key_matches(key, n):
                    if check_vram and vram and spec.vram_gb and abs(vram - spec.vram_gb) > 4:
                        continue
                    return bw, "vendor spec sheet"
        return None, ""
    return None, ""


# --------------------------------------------------------------------------
# Trained model (optional, loaded if scripts/train_speed_model.py produced it)
# --------------------------------------------------------------------------

_MODEL_JSON_PATH = _ROOT / "model" / "speed_model.json"


@lru_cache(maxsize=1)
def _trained_model():
    # Prefer XGBoost's NATIVE format: zero extra deps at runtime (the skops
    # artifact exists for the Hub, but its loading chain dragged in unrelated
    # imports on the Space).
    if _MODEL_JSON_PATH.exists():
        try:
            from xgboost import XGBRegressor
            model = XGBRegressor()
            model.load_model(_MODEL_JSON_PATH)
            print(f"[FitCheck] speed predictor loaded from {_MODEL_JSON_PATH.name}", flush=True)
            return model
        except Exception as e:  # noqa: BLE001
            import sys
            print(f"[FitCheck] WARNING: {_MODEL_JSON_PATH.name} exists but failed "
                  f"to load ({e!r}) — trying the skops artifact",
                  file=sys.stderr, flush=True)
    if not _MODEL_PATH.exists():
        return None
    try:
        from skops.io import load as skops_load
        # skops only loads explicitly-trusted types — exactly these two, which
        # scripts/train_speed_model.py produces. Anything else is refused.
        model = skops_load(_MODEL_PATH, trusted=[
            "xgboost.core.Booster", "xgboost.sklearn.XGBRegressor",
        ])
        print(f"[FitCheck] speed predictor loaded from {_MODEL_PATH.name}", flush=True)
        return model
    except Exception as e:  # noqa: BLE001
        # The file exists but won't load — say so loudly (a silent fallback
        # here would hide a broken deploy behind plausible roofline numbers).
        import sys
        print(f"[FitCheck] WARNING: {_MODEL_PATH.name} exists but failed to "
              f"load ({e!r}) — falling back to the labelled roofline estimate",
              file=sys.stderr, flush=True)
        return None


_METRICS_PATH = _ROOT / "model" / "metrics.json"


@lru_cache(maxsize=1)
def _envelope() -> dict:
    """The region of feature space the training data actually covered.

    Decision trees cannot extrapolate: outside what they saw, they clamp to
    the nearest seen value and quietly give wrong answers (e.g. a 32B model
    gets a 14B's speed). The roofline DOES extrapolate — it's physics. So the
    trained model only answers inside its measured envelope; outside it, the
    labelled analytical estimate takes over. Bounds come from metrics.json
    when the training script recorded them, else conservative defaults
    matching the LocalScore grid (<=14B Q4 models, consumer hardware).
    """
    env = {"bytes_gb": (0.8, 10.0), "eff_bw": (30.0, 1900.0)}
    try:
        rec = json.loads(_METRICS_PATH.read_text(encoding="utf-8")).get("envelope")
        if rec:
            env.update({k: tuple(v) for k, v in rec.items()})
    except OSError:
        pass
    return env


def _in_envelope(eff_bw: float, bytes_gb: float) -> bool:
    env = _envelope()
    return (env["bytes_gb"][0] <= bytes_gb <= env["bytes_gb"][1]
            and env["eff_bw"][0] <= eff_bw <= env["eff_bw"][1])


# --------------------------------------------------------------------------
# Prediction
# --------------------------------------------------------------------------

def predict_decode_tps(
    *,
    bandwidth_gbs: float,
    weights_gb: float,
    kv_gb: float = 0.0,
    active_fraction: float = 1.0,
    offload_fraction: float = 0.0,
) -> dict:
    """Predict decode tokens/sec.

    active_fraction: MoE models only read their active experts per token.
    offload_fraction: share of the model living in system RAM (0 = all on GPU).
    """
    # Bytes read per generated token: the (active) weights + the KV cache.
    bytes_gb = max(weights_gb * active_fraction + kv_gb, 0.05)
    if active_fraction < 0.9:
        # MoE conservatism: expert routing scatters reads across the full
        # weight file, so real MoE decode lands well under the active-bytes
        # ideal. 1.5x is a deliberate under-promise until measured data
        # corrects it (community MoE numbers run ~50-70% of ideal).
        bytes_gb *= 1.5

    eff_bw = bandwidth_gbs
    if offload_fraction > 0:
        f = min(max(offload_fraction, 0.0), 1.0)
        eff_bw = 1.0 / ((1.0 - f) / bandwidth_gbs + f / _RAM_BW_GBS)

    # The trained model only ever saw dense, fully-resident measurements
    # (active_fraction == 1, offload_fraction == 0 — LocalScore's grid). MoE
    # routing and GPU->RAM offload are corrected ANALYTICALLY above (the *1.5
    # MoE factor and the blended eff_bw); feeding those regimes to a tree that
    # never saw them is extrapolation, so we drop to the roofline there instead.
    learned_ok = active_fraction >= 0.999 and offload_fraction <= 0.0

    model = _trained_model()
    if model is not None and learned_ok and _in_envelope(eff_bw, bytes_gb):
        try:
            import numpy as np
            x = np.array([[eff_bw, bytes_gb, weights_gb, kv_gb,
                           active_fraction, offload_fraction,
                           eff_bw / bytes_gb]])
            tps = float(model.predict(x)[0])
            return {"tps": round(tps, 1),
                    "lo": round(tps * 0.8, 1), "hi": round(tps * 1.2, 1),
                    "bytes_gb": round(bytes_gb, 2), "eff_bw": round(eff_bw, 1),
                    "method": "measured-model",
                    "note": ("predicted by a model trained on real community "
                             "measurements (LocalScore), LLM-Pilot methodology")}
        except Exception:  # noqa: BLE001 — fall through to roofline
            pass

    # Per-token time = bandwidth-bound read time + fixed overhead. The overhead
    # term keeps tiny models honest (they plateau, they don't hit 1000s tok/s).
    def _tps(eff):
        return 1.0 / (bytes_gb / (eff_bw * eff) + _TOKEN_OVERHEAD_S)
    note = ("analytical estimate: decode speed is memory-bandwidth-bound "
            "(bandwidth divided by bytes read per token), with a fixed per-token "
            "overhead so small models plateau rather than scaling without limit")
    if model is not None:
        note += (" — this configuration is outside the measured data's range, "
                 "so the physics formula answers instead of the trained model")
    return {"tps": round(_tps(_EFF_MID), 1),
            "lo": round(_tps(_EFF_LO), 1), "hi": round(_tps(_EFF_HI), 1),
            "bytes_gb": round(bytes_gb, 2), "eff_bw": round(eff_bw, 1),
            "method": "roofline", "note": note}


def feel_text(pred: dict) -> str:
    """One honest, plain-English line from a prediction."""
    tps = pred["tps"]
    lo, hi = pred["lo"], pred["hi"]
    if tps >= _READING_TPS * 4:
        speed_word = "much faster than you read"
    elif tps >= _READING_TPS * 1.5:
        speed_word = "faster than you read"
    elif tps >= _READING_TPS * 0.7:
        speed_word = "about reading speed"
    else:
        speed_word = "slower than reading — fine for short tasks"
    return f"~{tps:g} tok/s (likely {lo:g}-{hi:g}) — {speed_word}"


# ==========================================================================
# Compute-bound speed (vision detection + diffusion). A DIFFERENT axis from
# LLM decode: these are limited by tensor-core FLOPs, not memory bandwidth.
#
#   latency_s = GFLOPs_per_forward / (peak_TFLOPS * 1000 * utilisation)
#
# `utilisation` (Model FLOPs Utilisation) is the fudge factor and is NOT
# universal — it's calibrated per family at batch=1 (FitCheck's case: one
# image / one prompt). Bands below are fit on the SAME basis as the TFLOPS
# column in gpu_specs.json (FP16-accumulate dense), so they transfer across
# GPUs:
#   detection: Ultralytics YOLO11/26 official T4 TensorRT latencies vs their
#     published GFLOPs (T4 peak 65 TFLOPS). util rises with model size
#     (~0.05 nano -> ~0.27 x-large) — a power fit, not a constant.
#   diffusion: real SDXL it/s on A100/4090/3090 -> util 0.124/0.138/0.152 on
#     this basis (mid ~0.14). Per-UNet/DiT step, batch=1.
#
# Honesty rules baked in (SPEED-BRICK-RESEARCH.md §9): emit a BAND not a point;
# FLOPs is a latency LOWER bound (FPS UPPER bound) and a documented imperfect
# GPU proxy, so tiny models on fast GPUs are often launch-bound BELOW this
# ceiling — the caption says so. Only NVIDIA GPUs (which have a calibrated
# tflops figure) and only the two calibrated families get a number; anything
# else returns None and the caller shows a non-speed verdict.
# ==========================================================================

# Diffusion utilisation band (batch=1, per-step, FP16-acc basis).
_UTIL_DIFF_LO, _UTIL_DIFF_MID, _UTIL_DIFF_HI = 0.11, 0.14, 0.17
# Above this FPS the exact number is meaningless for a person (and the model is
# launch-bound, not compute-bound), so we say "N+ FPS" instead.
_FPS_DISPLAY_CAP = 300.0


@lru_cache(maxsize=1)
def _tflops_index() -> tuple:
    idx = []
    for name, d in _specs()["gpus"].items():
        if d.get("tflops"):
            idx.append((_norm(name), float(d["tflops"])))
    idx.sort(key=lambda t: -len(t[0]))   # longest first: '4080 super' beats '4080'
    return tuple(idx)


def tflops_for_spec(spec, gpu_label: str = "") -> tuple[float | None, str]:
    """(peak FP16 tensor TFLOPS, source-note) for a machine, or (None, '').

    Only NVIDIA cards carry a calibrated tflops figure; Apple/AMD/Intel return
    None on purpose (no calibrated vision utilisation band exists for them yet,
    so the engine says 'out of scope' rather than guessing)."""
    if getattr(spec, "is_apple_silicon", False):
        return None, ""
    if spec.gpu_vendor == "nvidia" and spec.vram_gb > 0:
        n = _norm(gpu_label or spec.gpu_label)
        for key, tf in _tflops_index():
            if key and _key_matches(key, n):
                return tf, "vendor tensor-core spec"
    return None, ""


def _util_det(gflops: float) -> float:
    """Detection MFU vs model size. Log-linear fit to the YOLO11/26 T4 anchors
    (GFLOPs -> measured util): 5.4->0.05, 20.7->0.13, 68->0.22, 87->0.22,
    194->0.27. Bigger models keep the tensor cores busier per launch."""
    import math
    return min(0.30, max(0.05, 0.139 * math.log10(max(gflops, 1.0)) - 0.053))


def predict_detection_fps(gflops: float, tflops: float) -> dict:
    """Compute-roofline frames/sec for a single-shot detector at batch 1."""
    util = _util_det(gflops)

    def _fps(u):
        return tflops * 1000.0 * u / gflops

    mid = _fps(util)
    lo = _fps(util * 0.75)      # band reflects batch=1 utilisation spread
    hi = _fps(util * 1.25)
    return {
        "kind": "detection",
        "fps": round(mid, 1), "lo": round(lo, 1), "hi": round(hi, 1),
        "tflops": tflops, "util": round(util, 3),
        "realtime": lo >= 30.0,
        "method": "compute-roofline",
        "note": ("compute-roofline estimate (tensor FLOPs / model FLOPs, "
                 "batch 1, FP16) calibrated on Ultralytics' T4 latencies. This "
                 "is a ceiling: small models on a fast GPU are often limited by "
                 "per-frame launch overhead below it."),
    }


def predict_diffusion_time(gflops_step: float, steps: int, overhead_gflops: float,
                           tflops: float) -> dict:
    """Compute-roofline seconds-per-image for a diffusion model at batch 1."""
    total_gflops = steps * gflops_step + overhead_gflops

    def _sec(u):
        return total_gflops / (tflops * 1000.0 * u)

    mid = _sec(_UTIL_DIFF_MID)
    fast = _sec(_UTIL_DIFF_HI)      # higher util -> less time
    slow = _sec(_UTIL_DIFF_LO)
    return {
        "kind": "diffusion",
        "seconds": round(mid, 1), "sec_lo": round(fast, 1), "sec_hi": round(slow, 1),
        "it_s": round(steps / mid, 1), "steps": steps,
        "tflops": tflops, "util": _UTIL_DIFF_MID,
        "method": "compute-roofline",
        "note": ("compute-roofline estimate (tensor FLOPs / model FLOPs, "
                 "batch 1, FP16) calibrated on measured SDXL throughput. Assumes "
                 "an unoptimised baseline; step-skipping or caching is faster."),
    }


def predict_compute_speed(entry: dict, spec, gpu_label: str = "") -> dict | None:
    """Dispatch a catalogue entry to the right compute-roofline estimate, or
    None when it's out of scope (non-NVIDIA GPU, or a family with no calibrated
    utilisation band). The `compute_kind` tag is set in the catalogue by
    scripts/add_model_gflops.py for the calibrated entries only."""
    kind = entry.get("compute_kind")
    if kind not in ("detection", "diffusion"):
        return None
    tflops, src = tflops_for_spec(spec, gpu_label)
    if not tflops:
        return None
    if kind == "detection" and entry.get("gflops"):
        pred = predict_detection_fps(float(entry["gflops"]), tflops)
    elif kind == "diffusion" and entry.get("gflops_step"):
        pred = predict_diffusion_time(
            float(entry["gflops_step"]), int(entry.get("steps_default", 30)),
            float(entry.get("gflops_overhead", 0.0)), tflops)
    else:
        return None
    pred["tflops_source"] = src
    pred["proxy"] = bool(entry.get("gflops_proxy"))
    return pred


def compute_feel_text(pred: dict) -> str:
    """One honest, plain-English line from a compute-roofline prediction."""
    if pred["kind"] == "detection":
        fps, lo, hi = pred["fps"], pred["lo"], pred["hi"]
        if lo >= _FPS_DISPLAY_CAP:
            return f"{int(_FPS_DISPLAY_CAP)}+ FPS — far beyond real-time (compute ceiling)"
        rt = "real-time" if pred["realtime"] else "near real-time" if hi >= 30 else "below real-time"
        return f"~{fps:g} FPS (likely {lo:g}-{hi:g}) — {rt} (compute ceiling, batch 1)"
    # diffusion
    s, lo, hi = pred["seconds"], pred["sec_lo"], pred["sec_hi"]
    return f"~{s:g}s per image (likely {lo:g}-{hi:g}s) at {pred['steps']} steps"