"""Device management utilities for CUDA and MPS support."""

import torch
import torch.nn as nn


def get_device(config_device: str = "auto") -> torch.device:
    """Get the best available device.

    Args:
        config_device: Device configuration ("auto", "cuda", "mps", "cpu").

    Returns:
        torch.device for computation.
    """
    if config_device != "auto":
        device = torch.device(config_device)
        print(f"Using configured device: {device}")
        return device

    if torch.cuda.is_available():
        device = torch.device("cuda")
        print(f"Using CUDA: {torch.cuda.get_device_name()}")
        print(f"  Memory: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB")
        return device

    if torch.backends.mps.is_available():
        device = torch.device("mps")
        print("Using Apple MPS (Metal Performance Shaders)")
        return device

    print("Using CPU")
    return torch.device("cpu")


def setup_for_device(model: nn.Module, device: torch.device) -> nn.Module:
    """Setup model for specific device with optimizations.

    Args:
        model: PyTorch model.
        device: Target device.

    Returns:
        Model moved to device with optimizations applied.
    """
    model = model.to(device)

    if device.type == "cuda":
        # Enable TF32 for faster computation on Ampere+ GPUs
        torch.backends.cuda.matmul.allow_tf32 = True
        torch.backends.cudnn.allow_tf32 = True

        # Enable cudnn benchmark for consistent input sizes
        torch.backends.cudnn.benchmark = True
        print("  CUDA optimizations enabled (TF32, cudnn benchmark)")

    elif device.type == "mps":
        # MPS-specific settings
        print("  MPS device configured")

    return model


def get_device_info(device: torch.device) -> dict:
    """Get information about the device.

    Args:
        device: The device to get info for.

    Returns:
        Dictionary with device information.
    """
    info = {
        "device_type": device.type,
        "device_index": device.index,
    }

    if device.type == "cuda":
        info.update(
            {
                "cuda_device_name": torch.cuda.get_device_name(device),
                "cuda_capability": torch.cuda.get_device_capability(device),
                "cuda_memory_total": torch.cuda.get_device_properties(device).total_memory,
                "cuda_memory_allocated": torch.cuda.memory_allocated(device),
                "cuda_memory_cached": torch.cuda.memory_reserved(device),
            }
        )
    elif device.type == "mps":
        info["mps_available"] = torch.backends.mps.is_available()
        info["mps_built"] = torch.backends.mps.is_built()

    return info
