import itertools

import os

from typing import Union

import torch

def get_default_device():

    if torch.cuda.is_available():

        return get_most_free_gpu()

    elif torch.backends.mps.is_available():

        # Not all operations implemented in MPS yet

        use_mps = os.environ.get("PYTORCH_ENABLE_MPS_FALLBACK", "0") == "1"

        if use_mps:

            return torch.device('mps')

        else:

            return torch.device('cpu')

    else:

        return torch.device('cpu')

def get_available_devices(num_devices=None, max_devices=None):

    device = get_default_device()

    if device.type == "cuda":

        num_gpus = torch.cuda.device_count()

        if max_devices is not None:

            num_gpus = min(num_gpus, max_devices)

        gpu_list = [get_device(i) for i in range(num_gpus)]

        if num_devices is not None:

            cycle_gpu_list = itertools.cycle(gpu_list)

            gpu_list = [next(cycle_gpu_list) for _ in range(num_devices)]

        return gpu_list

    else:

        num_devices = num_devices if num_devices else torch.multiprocessing.cpu_count()

        num_devices = min(num_devices, max_devices) if max_devices is not None else num_devices

        return [device]*num_devices

def get_device(gpu_id: int):

    if gpu_id is not None and torch.cuda.is_available() and gpu_id < torch.cuda.device_count():

        return torch.device(f'cuda:{gpu_id}')

    else:

        return None

def get_device_mem_info(device: Union[int, torch.device]):

    if not torch.cuda.is_available():

        return None

    free_mem, total_mem = torch.cuda.mem_get_info(device=device)

    return free_mem, total_mem

def get_most_free_gpu():

    """

    Get the GPU with the most free memory

    If system has no GPUs (or CUDA not available), return None

    """

    if not torch.cuda.is_available():

        return None

    num_gpus = torch.cuda.device_count()

    if num_gpus == 0:

        return None

    most_free_idx, most_free_val = -1, -1

    for i in range(num_gpus):

        free_mem, total_mem = get_device_mem_info(i)

        if free_mem > most_free_val:

            most_free_idx, most_free_val = i, free_mem

    return torch.device(f'cuda:{most_free_idx}')