from __future__ import annotations

from .models import AcceleratorProfile, ModelProfile


MODELS: dict[str, ModelProfile] = {
    "Llama-3.1-8B": ModelProfile(
        name="Llama-3.1-8B",
        params_b=8.03,
        layers=32,
        hidden_size=4096,
        attention_heads=32,
        kv_heads=8,
    ),
    "Mistral-7B-v0.3": ModelProfile(
        name="Mistral-7B-v0.3",
        params_b=7.25,
        layers=32,
        hidden_size=4096,
        attention_heads=32,
        kv_heads=8,
    ),
    "Qwen2.5-3B": ModelProfile(
        name="Qwen2.5-3B",
        params_b=3.09,
        layers=36,
        hidden_size=2048,
        attention_heads=16,
        kv_heads=2,
    ),
}


ACCELERATORS: dict[str, AcceleratorProfile] = {
    "L4": AcceleratorProfile(
        name="NVIDIA L4",
        vram_gb=24.0,
        peak_tflops_fp16=121.0,
        bandwidth_gbps=300.0,
        compute_efficiency=0.40,
        bandwidth_efficiency=0.72,
    ),
    "A10G": AcceleratorProfile(
        name="NVIDIA A10G",
        vram_gb=24.0,
        peak_tflops_fp16=125.0,
        bandwidth_gbps=600.0,
        compute_efficiency=0.40,
        bandwidth_efficiency=0.70,
    ),
    "A100-40GB": AcceleratorProfile(
        name="NVIDIA A100 40GB",
        vram_gb=40.0,
        peak_tflops_fp16=312.0,
        bandwidth_gbps=1555.0,
        compute_efficiency=0.47,
        bandwidth_efficiency=0.76,
    ),
}


QUANTIZATION_BYTES = {
    "fp16": 2.0,
    "int8": 1.0,
    "int4": 0.5,
}

# Compute dequantization / packing overheads are deliberately conservative analytical
# modifiers, not empirical benchmark claims.
QUANTIZATION_COMPUTE_MULTIPLIER = {
    "fp16": 1.00,
    "int8": 1.07,
    "int4": 1.16,
}


def get_model(name: str) -> ModelProfile:
    try:
        return MODELS[name]
    except KeyError as exc:
        raise ValueError(f"Unknown model profile: {name}") from exc


def get_accelerator(name: str) -> AcceleratorProfile:
    try:
        return ACCELERATORS[name]
    except KeyError as exc:
        raise ValueError(f"Unknown accelerator profile: {name}") from exc