"""
Extrahierte Mathematik aus dem Odysseus Projekt zur VRAM/RAM Berechnung.
Abgestimmt auf APUs mit Unified Memory (Bosgame M5 / Strix Halo).
"""

import re

# Annahme: Bytes per Parameter für GGUF Quants
QUANT_BYTES_PER_PARAM = {
    "Q2_K": 0.35,
    "Q3_K_S": 0.38,
    "Q3_K_M": 0.42,
    "Q3_K_L": 0.45,
    "Q4_0": 0.50,
    "Q4_1": 0.55,
    "Q4_K_S": 0.50,
    "Q4_K_M": 0.55,
    "Q5_0": 0.62,
    "Q5_1": 0.68,
    "Q5_K_S": 0.62,
    "Q5_K_M": 0.65,
    "Q6_K": 0.75,
    "Q8_0": 1.00,
    "F16": 2.00,
    "BF16": 2.00,
}

def estimate_memory_gb(params_b: float, quant: str, ctx: int) -> float:
    """Berechnet den geschätzten Speicherbedarf in GB (Gewichte + Kontext)."""
    # Wenn unbekanntes Format, nimm sicherheitshalber Q5_K_M (0.65)
    bpp = QUANT_BYTES_PER_PARAM.get(quant.upper(), 0.65)
    weights = params_b * bpp
    
    # Heuristik für Context-RAM: 8k Context bei 7B Parametern frisst ca. 0.8 GB
    context_vram = (ctx / 8192) * (max(params_b, 7) / 7) * 0.8
    
    return weights + context_vram

def extract_active_params_b(name: str) -> float | None:
    """Aktive Parameter bei MoE-Modellen aus dem Namen (z.B. '30B-A3B' → 3.0).
    Gibt None zurück für Dense-Modelle (kein MoE-Suffix erkannt)."""
    m = re.search(r'(?<![a-zA-Z])a(\d+(?:\.\d+)?)b\b', name.lower())
    return float(m.group(1)) if m else None


def estimate_speed(req_gb: float, sys_ram_gb: float, moe_active_ratio: float = 1.0) -> float:
    """Berechnet die geschätzte Tokens/s basierend auf der 273 GB/s Bandbreite der APU.
    moe_active_ratio = aktive_params / gesamt_params; < 1 bei MoE (z.B. 0.1 für 30B-A3B)."""
    bw = 273 if sys_ram_gb > 8 else 70
    if req_gb <= 0:
        return 0.0
    raw_tps = (bw / req_gb) * 0.55
    # MoE-Boost: nur ein Bruchteil der Gewichte wird pro Token aktiviert →
    # effektiv höhere Durchsatzrate als ein Dense-Modell gleicher Gesamtgröße.
    # sqrt-Dämpfung: Routing-Overhead + VRAM-Traffic für alle Experten bleibt.
    if moe_active_ratio < 0.8:
        raw_tps *= (1.0 / moe_active_ratio) ** 0.5
    return raw_tps


def evaluate_fit(params_b: float, quant: str, ctx: int, sys_ram_gb: float,
                 name: str = "") -> dict:
    """Berechnet den Fit für ein System mit Shared Memory (APU).
    name wird für MoE-Erkennung genutzt (optional)."""
    req_gb = estimate_memory_gb(params_b, quant, ctx)
    active_b = extract_active_params_b(name) if name else None
    moe_ratio = (active_b / params_b) if (active_b and params_b > 0) else 1.0
    tps = estimate_speed(req_gb, sys_ram_gb, moe_ratio)
    
    # Das OS und andere Prozesse brauchen RAM. Wir lassen 4GB Puffer.
    usable_ram = max(sys_ram_gb - 4.0, 0)
    
    if req_gb > usable_ram:
        fit_level = "too_tight"
        text = "Zu groß (OOM)"
    elif req_gb > usable_ram * 0.8:
        fit_level = "marginal"
        text = "Könnte knapp werden"
    else:
        fit_level = "perfect"
        text = "Passt perfekt"
        
    return {
        "level": fit_level,
        "text": text,
        "req_gb": round(req_gb, 1),
        "tps": round(tps, 0)
    }


def extract_params_b(name: str) -> float:
    """Parametergröße (Mrd.) aus Repo-/Dateiname. 8x7B (MoE) -> 56."""
    moe = re.search(r"(\d+)x(\d+(?:\.\d+)?)[bB]", name)
    if moe:
        return float(moe.group(1)) * float(moe.group(2))
    m = re.search(r"(\d+(?:\.\d+)?)[bB](?![a-zA-Z])", name)
    if m:
        return float(m.group(1))
    return 7.0


# "Schöne" Kontextstufen, die UIs/Modelle gern mögen.
_NICE_CTX = [2048, 4096, 8192, 16384, 32768, 49152, 65536, 98304, 131072]


def max_ctx_for(params_b: float, quant: str, sys_ram_gb: float) -> int:
    """Größter 'schöner' Kontext, der komfortabel passt (80% des nutzbaren RAM, 4 GB OS-Puffer).
    Umkehrung von estimate_memory_gb: löst die Kontext-Heuristik nach ctx auf."""
    bpp = QUANT_BYTES_PER_PARAM.get(quant.upper(), 0.65)
    weights = params_b * bpp
    usable = max(sys_ram_gb - 4.0, 0) * 0.8           # gleicher Komfort wie evaluate_fit "perfect"
    ctx_budget = usable - weights                      # GB, die für den KV-Cache übrig sind
    if ctx_budget <= 0:
        return 2048                                    # Modell selbst schon knapp -> Minimal-Kontext
    per_8k = (max(params_b, 7) / 7) * 0.8              # GB pro 8k Kontext (aus der Heuristik)
    raw_ctx = (ctx_budget / per_8k) * 8192
    best = _NICE_CTX[0]
    for c in _NICE_CTX:
        if c <= raw_ctx:
            best = c
    return best


def recommend_ctx(params_b: float, quant: str, sys_ram_gb: float) -> dict:
    """Empfohlener Kontext + Klartext-Begründung (für die UI)."""
    ctx = max_ctx_for(params_b, quant, sys_ram_gb)
    k = ctx // 1024
    return {"ctx": ctx, "k": k,
            "note": f"Bis ~{k}k Kontext passt komfortabel auf deine Hardware ({round(sys_ram_gb)} GB)."}