"""Hardware the fit estimate plans for.

Browsers cannot read VRAM or unified-memory size, so the user picks a device (or
enters numbers). Discrete GPUs lose some memory to the driver/context; unified-memory
systems let the GPU use only part of RAM by default.
"""

from __future__ import annotations

from dataclasses import dataclass

# name -> VRAM in GB (as sold), grouped the way the page lists them. Names already stored
# in the dataset (e.g. "NVIDIA RTX 3060 12GB", "Intel Arc A770 16GB") must not change.
GPU_GROUPS: list[tuple[str, dict[str, float]]] = [
    ("NVIDIA RTX 50", {
        "NVIDIA RTX 5050": 8, "NVIDIA RTX 5060": 8, "NVIDIA RTX 5060 Ti 8GB": 8, "NVIDIA RTX 5060 Ti 16GB": 16,
        "NVIDIA RTX 5070": 12, "NVIDIA RTX 5070 Ti": 16, "NVIDIA RTX 5080": 16, "NVIDIA RTX 5090": 32}),
    ("NVIDIA RTX 40", {
        "NVIDIA RTX 4060": 8, "NVIDIA RTX 4060 Ti 8GB": 8, "NVIDIA RTX 4060 Ti 16GB": 16, "NVIDIA RTX 4070": 12,
        "NVIDIA RTX 4070 Super": 12, "NVIDIA RTX 4070 Ti": 12, "NVIDIA RTX 4070 Ti Super": 16, "NVIDIA RTX 4080": 16,
        "NVIDIA RTX 4080 Super": 16, "NVIDIA RTX 4090": 24, "NVIDIA RTX 4090 48GB (modded)": 48}),
    ("NVIDIA RTX 30", {
        "NVIDIA RTX 3050 8GB": 8, "NVIDIA RTX 3060 8GB": 8, "NVIDIA RTX 3060 12GB": 12, "NVIDIA RTX 3060 Ti": 8,
        "NVIDIA RTX 3070": 8, "NVIDIA RTX 3070 Ti": 8, "NVIDIA RTX 3080 10GB": 10, "NVIDIA RTX 3080 12GB": 12,
        "NVIDIA RTX 3080 Ti": 12, "NVIDIA RTX 3090": 24, "NVIDIA RTX 3090 Ti": 24}),
    ("NVIDIA RTX 20 / GTX 16", {
        "NVIDIA GTX 1650": 4, "NVIDIA GTX 1660": 6, "NVIDIA GTX 1660 Super": 6, "NVIDIA GTX 1660 Ti": 6,
        "NVIDIA RTX 2060 6GB": 6, "NVIDIA RTX 2060 12GB": 12, "NVIDIA RTX 2060 Super": 8, "NVIDIA RTX 2070": 8,
        "NVIDIA RTX 2070 Super": 8, "NVIDIA RTX 2080": 8, "NVIDIA RTX 2080 Super": 8, "NVIDIA RTX 2080 Ti": 11,
        "NVIDIA RTX 2080 Ti 22GB (modded)": 22, "NVIDIA Titan RTX": 24}),
    ("NVIDIA GTX 10", {
        "NVIDIA GTX 1050 Ti": 4, "NVIDIA GTX 1060 3GB": 3, "NVIDIA GTX 1060 6GB": 6, "NVIDIA GTX 1070": 8,
        "NVIDIA GTX 1070 Ti": 8, "NVIDIA GTX 1080": 8, "NVIDIA GTX 1080 Ti": 11, "NVIDIA Titan X (Pascal)": 12,
        "NVIDIA Titan Xp": 12}),
    ("NVIDIA laptop", {
        "NVIDIA RTX 5050 Laptop": 8, "NVIDIA RTX 5060 Laptop": 8, "NVIDIA RTX 5070 Laptop": 8,
        "NVIDIA RTX 5070 Ti Laptop": 12, "NVIDIA RTX 5080 Laptop": 16, "NVIDIA RTX 5090 Laptop": 24,
        "NVIDIA RTX 4050 Laptop": 6, "NVIDIA RTX 4060 Laptop": 8, "NVIDIA RTX 4070 Laptop": 8,
        "NVIDIA RTX 4080 Laptop": 12, "NVIDIA RTX 4090 Laptop": 16,
        "NVIDIA RTX 3050 Laptop": 4, "NVIDIA RTX 3050 Ti Laptop": 4, "NVIDIA RTX 3060 Laptop": 6,
        "NVIDIA RTX 3070 Laptop": 8, "NVIDIA RTX 3070 Ti Laptop": 8, "NVIDIA RTX 3080 Laptop 8GB": 8,
        "NVIDIA RTX 3080 Laptop 16GB": 16, "NVIDIA RTX 3080 Ti Laptop": 16,
        "NVIDIA RTX 2060 Laptop": 6, "NVIDIA RTX 2070 Laptop": 8, "NVIDIA RTX 2080 Laptop": 8,
        "NVIDIA GTX 1650 Laptop": 4, "NVIDIA GTX 1660 Ti Laptop": 6}),
    ("NVIDIA workstation", {
        "NVIDIA Quadro P6000": 24, "NVIDIA Quadro RTX 8000": 48, "NVIDIA RTX A2000": 12, "NVIDIA RTX A4000": 16,
        "NVIDIA RTX A4500": 20, "NVIDIA RTX A5000": 24, "NVIDIA RTX A6000": 48, "NVIDIA RTX 2000 Ada": 16,
        "NVIDIA RTX 4000 Ada": 20, "NVIDIA RTX 4500 Ada": 24, "NVIDIA RTX 5000 Ada": 32, "NVIDIA RTX 6000 Ada": 48,
        "NVIDIA RTX PRO 4000 Blackwell": 24, "NVIDIA RTX PRO 4500 Blackwell": 32, "NVIDIA RTX PRO 5000 Blackwell": 48,
        "NVIDIA RTX PRO 6000 Blackwell": 96}),
    ("NVIDIA datacenter", {
        "NVIDIA Tesla M40 24GB": 24, "NVIDIA Tesla P4": 8, "NVIDIA Tesla P40": 24, "NVIDIA Tesla P100": 16,
        "NVIDIA Tesla V100 16GB": 16, "NVIDIA Tesla V100 32GB": 32, "NVIDIA T4": 16, "NVIDIA A2": 16,
        "NVIDIA A10": 24, "NVIDIA A30": 24, "NVIDIA A40": 48, "NVIDIA L4": 24, "NVIDIA L40S": 48,
        "NVIDIA A100 40GB": 40, "NVIDIA A100 80GB": 80, "NVIDIA H100": 80, "NVIDIA H200": 141, "NVIDIA B200": 180}),
    ("AMD Radeon", {
        "AMD RX 580 8GB": 8, "AMD RX 5700 XT": 8, "AMD Radeon VII": 16, "AMD RX 6600": 8, "AMD RX 6600 XT": 8,
        "AMD RX 6700 XT": 12, "AMD RX 6750 XT": 12, "AMD RX 6800": 16, "AMD RX 6800 XT": 16, "AMD RX 6900 XT": 16,
        "AMD RX 6950 XT": 16, "AMD RX 7600": 8, "AMD RX 7600 XT": 16, "AMD RX 7700 XT": 12, "AMD RX 7800 XT": 16,
        "AMD RX 7900 GRE": 16, "AMD RX 7900 XT": 20, "AMD RX 7900 XTX": 24, "AMD RX 9060 XT 8GB": 8,
        "AMD RX 9060 XT 16GB": 16, "AMD RX 9070": 16, "AMD RX 9070 XT": 16}),
    ("AMD workstation", {
        "AMD Radeon Pro W7800": 32, "AMD Radeon Pro W7900": 48, "AMD Radeon AI PRO R9700": 32,
        "AMD Instinct MI50 16GB": 16, "AMD Instinct MI50 32GB": 32, "AMD Instinct MI60": 32}),
    ("Intel Arc", {
        "Intel Arc A380": 6, "Intel Arc A580": 8, "Intel Arc A750": 8, "Intel Arc A770 8GB": 8,
        "Intel Arc A770 16GB": 16, "Intel Arc B570": 10, "Intel Arc B580": 12, "Intel Arc Pro B50": 16,
        "Intel Arc Pro B60": 24}),
]
GPUS: dict[str, float] = {name: vram for _, cards in GPU_GROUPS for name, vram in cards.items()}

# Unified-memory systems: name -> (total RAM choices, share the GPU can use by default)
UNIFIED = {
    "Apple Silicon Mac": ([8, 16, 18, 24, 32, 36, 48, 64, 96, 128, 192, 256, 512], None),
    "AMD Ryzen AI Max+ 395 (Strix Halo)": ([32, 64, 96, 128], 0.9),
    "NVIDIA DGX Spark": ([128], 0.9),
    # Ex-mining board with a cut-down PS5 APU: 16 GB GDDR6 shared, 12 GB to the GPU is the
    # baseline the community settled on, and llama.cpp runs on Vulkan (gfx1013 has no ROCm).
    "AMD BC-250 (PS5 APU)": ([16], 0.75),
}

UNIFIED_NOTES = {
    "AMD BC-250 (PS5 APU)": "Build llama.cpp with Vulkan; ROCm doesn't support this GPU. Set VRAM to 12 GB in "
                            "the BIOS and raise ttm.pages_limit, or the GPU can only reach about half the memory.",
}

RAM_CHOICES = [8, 16, 24, 32, 48, 64, 96, 128, 192, 256, 384, 512, 768, 1024]
CONTEXTS = [4096, 8192, 16384, 32768, 65536, 131072, 262144]
KV_TYPES = {"f16": 2.0, "q8_0": 1.0625, "q4_0": 0.5625}

GPU_CONTEXT_OVERHEAD_GB = 0.5   # CUDA/ROCm context + backend workspace, per GPU
HOST_OVERHEAD_GB = 1.0          # llama.cpp process, OS page cache headroom


@dataclass(frozen=True)
class Hardware:
    kind: str                      # "gpu" | "unified" | "cpu"
    gpus: tuple[float, ...] = ()   # VRAM per GPU, GB
    ram_gb: float = 32.0           # system RAM (unified: total memory)
    gpu_share: float | None = None  # unified: share of memory the GPU may use

    @property
    def vram_gb(self) -> float:
        return sum(self.gpus)

    def usable_vram_gb(self) -> float:
        if self.kind == "unified":
            share = self.gpu_share if self.gpu_share else (0.67 if self.ram_gb <= 36 else 0.75)
            return self.ram_gb * share
        return max(sum(v - GPU_CONTEXT_OVERHEAD_GB for v in self.gpus), 0.0)

    def label(self) -> str:
        if self.kind == "unified":
            return f"{self.ram_gb:g} GB unified memory"
        if self.kind == "gpu":
            n = len(self.gpus)
            each = f"{self.gpus[0]:g} GB" if len(set(self.gpus)) == 1 else "+".join(f"{v:g}" for v in self.gpus) + " GB"
            return f"{n}x {each} GPU + {self.ram_gb:g} GB RAM" if n > 1 else f"{each} GPU + {self.ram_gb:g} GB RAM"
        return f"CPU with {self.ram_gb:g} GB RAM"


def presets() -> dict:
    return {"gpus": GPUS, "gpu_groups": [[label, list(cards)] for label, cards in GPU_GROUPS],
            "unified": {k: v[0] for k, v in UNIFIED.items()},
            "unified_shares": {k: v[1] for k, v in UNIFIED.items()}, "unified_notes": UNIFIED_NOTES,
            "ram": RAM_CHOICES, "contexts": CONTEXTS, "kv_types": list(KV_TYPES)}