"""Hardware detection and automatic runtime profile (CUDA / Apple MPS / CPU).

The *effective* batch size is part of the experiment (it changes results); the
micro-batch that fits in memory and the gradient accumulation to reach the
effective batch are hardware details chosen here, so the same config gives
comparable results on a laptop and on a rented GPU.
"""
from __future__ import annotations

import os
import platform
from dataclasses import asdict, dataclass

import torch


@dataclass
class RuntimeProfile:
    device: str
    device_name: str
    vram_gb: float | None
    amp_dtype: str | None  # "bf16" | "fp16" | None
    micro_batch: int
    grad_accum: int
    num_workers: int
    compile: bool
    threads: int

    def to_dict(self) -> dict:
        return asdict(self)


def detect_device(prefer: str = "auto") -> str:
    if prefer != "auto":
        return prefer
    if torch.cuda.is_available():
        return "cuda"
    if getattr(torch.backends, "mps", None) is not None and torch.backends.mps.is_available():
        return "mps"
    return "cpu"


def resolve_runtime(train_cfg: dict, model_params: int | None = None) -> RuntimeProfile:
    hw = train_cfg.get("hardware", {}) or {}
    device = detect_device(hw.get("device", "auto"))
    eff = int(train_cfg.get("batch_size", 256))
    vram = None
    name = platform.processor() or platform.machine()
    if device == "cuda":
        props = torch.cuda.get_device_properties(0)
        vram = props.total_memory / 1024**3
        name = props.name
        # activations dominate for these tiny models; ~1 GB per 4k samples of a ~5M model is a
        # conservative rule of thumb. `max_micro_batch` in the config always wins.
        per_gb = hw.get("samples_per_gb", 2048 if (model_params or 0) > 3e6 else 8192)
        cap = int(max(64, (vram - 1.5) * per_gb))
    elif device == "mps":
        cap = 1024
    else:
        cap = 512
    cap = min(cap, int(hw.get("max_micro_batch", cap)))
    micro = eff
    while micro > cap and micro % 2 == 0:
        micro //= 2
    accum = max(1, eff // micro)

    amp = hw.get("amp", "auto")
    if amp == "auto":
        if device == "cuda":
            amp = "bf16" if torch.cuda.is_bf16_supported() else "fp16"
        else:
            amp = None  # CPU autocast rarely helps small models; MPS fp16 is flaky for training
    elif amp in ("none", False):
        amp = None

    cpus = os.cpu_count() or 2
    workers = hw.get("num_workers", "auto")
    if workers == "auto":
        workers = min(8, max(0, cpus // 2 - 1)) if device != "cpu" else min(4, max(0, cpus // 4))
    threads = hw.get("threads", "auto")
    if threads == "auto":
        threads = max(1, cpus - int(workers)) if device == "cpu" else max(1, min(4, cpus))
    compile_ = hw.get("compile", "auto")
    if compile_ == "auto":
        compile_ = device == "cuda"  # launch-overhead bound tiny models gain most on GPU
    return RuntimeProfile(device, name, round(vram, 1) if vram else None, amp, micro, accum,
                          int(workers), bool(compile_), int(threads))


def autocast(profile: RuntimeProfile):
    if profile.amp_dtype is None:
        return torch.autocast(device_type="cpu", enabled=False)
    dtype = torch.bfloat16 if profile.amp_dtype == "bf16" else torch.float16
    return torch.autocast(device_type=profile.device, dtype=dtype)