"""
CPU Optimization Utilities

Provides CPU-specific performance tuning for inference:
- Thread count optimization (PyTorch OpenMP + inter-op)
- Best compute dtype detection (bfloat16 on modern CPUs, else float32)
- Available RAM measurement
- CPU capability reporting

Key insight for CPU inference at full accuracy:
  - Loading weights in float16 saves RAM
  - Computing in float32 (or bfloat16 on AVX-512 BF16 CPUs) keeps accuracy
  - Memory-mapped SafeTensors = OS-managed RAM, not Python-managed RAM
  - The OS page cache acts as a free, smart layer cache for us
"""

from __future__ import annotations

import gc
import logging
import os
import platform

import torch

logger = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# CPU capability detection
# ---------------------------------------------------------------------------

def supports_bf16() -> bool:
    """
    Check whether this CPU supports bfloat16 natively.

    Native BF16 support exists on:
    - Intel Sapphire Rapids (Xeon Gen 4, 2023+) via AMX
    - Intel Alder Lake / Raptor Lake (Core Gen 12/13) via AVX-512 BF16
    - AMD Zen 4 (EPYC Genoa, Ryzen 7000) via AVX-512 BF16
    - Apple M-series (always)

    On CPUs WITHOUT native BF16, PyTorch emulates it in software (slow).
    We only recommend BF16 when it's hardware-accelerated.
    """
    if platform.system() == "Darwin":
        # Apple Silicon always supports BF16 efficiently
        return True

    try:
        # PyTorch 2.1+ exposes CPU feature flags
        return torch.cpu.is_bf16_supported()
    except AttributeError:
        pass

    # Fallback: check /proc/cpuinfo for avx512bf16 flag (Linux)
    try:
        with open("/proc/cpuinfo") as f:
            for line in f:
                if line.startswith("flags") and "avx512_bf16" in line:
                    return True
    except OSError:
        pass

    return False


def get_compute_dtype() -> torch.dtype:
    """
    Return the optimal compute dtype for this CPU.

    - If BF16 is natively supported: torch.bfloat16 (same range as FP32,
      faster on modern CPUs, zero accuracy loss vs FP32 in most scenarios)
    - Otherwise: torch.float32 (universally safe, standard precision)

    NOTE: We never use float16 on CPU — float16 ALU is not native on x86,
    making it slower than float32 (emulated via float32 conversion).
    """
    if supports_bf16():
        logger.debug("CPU supports native BF16 — using bfloat16 for compute")
        return torch.bfloat16
    logger.debug("CPU does not have native BF16 — using float32 for compute")
    return torch.float32


def get_weight_load_dtype() -> torch.dtype:
    """
    Dtype to use when loading weights from disk.

    float16: Halves RAM usage vs float32, supported on all hardware for
             storage. PyTorch automatically upcasts during matmul.
    """
    return torch.float16


def setup_threads(n_compute: int | None = None, n_interop: int | None = None) -> None:
    """
    Configure PyTorch thread pools for optimal CPU inference.

    Args:
        n_compute: Threads for intra-op parallelism (matmul, conv, etc.).
                   Default: all physical CPU cores.
        n_interop: Threads for inter-op parallelism (parallel ops, I/O overlap).
                   Default: min(4, cpu_count // 2).
    """
    cpu_count = os.cpu_count() or 4

    if n_compute is None:
        # Use all available cores for compute
        n_compute = cpu_count

    if n_interop is None:
        # A few threads for async I/O overlap
        n_interop = min(4, max(1, cpu_count // 2))

    torch.set_num_threads(n_compute)
    torch.set_num_interop_threads(n_interop)

    logger.info(
        "CPU threads: %d compute, %d inter-op (of %d available)",
        n_compute, n_interop, cpu_count,
    )


def available_ram_bytes() -> int:
    """Return currently available system RAM in bytes."""
    try:
        import psutil
        return psutil.virtual_memory().available
    except ImportError:
        pass

    # Fallback: read /proc/meminfo on Linux
    try:
        with open("/proc/meminfo") as f:
            for line in f:
                if line.startswith("MemAvailable:"):
                    kb = int(line.split()[1])
                    return kb * 1024
    except OSError:
        pass

    # Conservative fallback: assume 2 GB available
    return 2 * 1024**3


def total_ram_bytes() -> int:
    """Return total system RAM in bytes."""
    try:
        import psutil
        return psutil.virtual_memory().total
    except ImportError:
        pass

    try:
        with open("/proc/meminfo") as f:
            for line in f:
                if line.startswith("MemTotal:"):
                    kb = int(line.split()[1])
                    return kb * 1024
    except OSError:
        pass

    return 4 * 1024**3  # assume 4 GB


def force_gc() -> int:
    """
    Aggressively free Python-managed memory.

    Returns bytes freed (approximate, based on tracked tensor sizes).
    This is important on CPU where PyTorch does not have a separate
    memory pool — it uses the system allocator directly.
    """
    collected = gc.collect()
    gc.collect()  # Two passes to handle reference cycles
    return collected


def cpu_info() -> dict:
    """Return a summary of CPU capabilities."""
    cpu_count = os.cpu_count() or 0
    return {
        "cpu_count": cpu_count,
        "platform": platform.processor() or platform.machine(),
        "supports_bf16": supports_bf16(),
        "compute_dtype": str(get_compute_dtype()),
        "total_ram_gb": round(total_ram_bytes() / 1e9, 1),
        "available_ram_gb": round(available_ram_bytes() / 1e9, 1),
        "torch_threads": torch.get_num_threads(),
    }


def print_cpu_info() -> None:
    """Print CPU capabilities to logger."""
    info = cpu_info()
    logger.info(
        "CPU info: %s | cores=%d | RAM=%.1f GB available / %.1f GB total | "
        "BF16=%s | compute_dtype=%s",
        info["platform"],
        info["cpu_count"],
        info["available_ram_gb"],
        info["total_ram_gb"],
        info["supports_bf16"],
        info["compute_dtype"],
    )
