import torch


def parallelize_model(model, device, num_gpus=None, gpu_ids=None, local_rank=-1):
    """Moves a model to the specified device (cpu or gpu/s)
       and implements data parallelism when multiple gpus are specified.
    Args:
        model (Module): A PyTorch model.
        device (torch.device): A PyTorch device.
        num_gpus (int): The number of GPUs to be used.
            If set to None, all available GPUs will be used.
            Defaults to None.
        gpu_ids (list): List of GPU IDs to be used.
            If None, the first num_gpus GPUs will be used.
            If not None, overrides num_gpus. if gpu_ids is an empty list
            or there is no valid gpu devices are specified,
            and device is "cuda", model will not be moved or parallelized.
            Defaults to None.
        local_rank (int): Local GPU ID within a node. Used in distributed environments.
            If not -1, num_gpus and gpu_ids are ignored.
            Defaults to -1.
    Returns:
        Module, DataParallel, DistributedDataParallel: A PyTorch Module or
            a DataParallel/DistributedDataParallel wrapper,
            when one or multiple gpus are used.
    """
    if not isinstance(device, torch.device):
        raise ValueError("device must be of type torch.device.")

    model_module = (
        model.module if hasattr(model, "module") else model
    )  # Take care of distributed/parallel training

    if local_rank != -1:
        model = torch.nn.parallel.DistributedDataParallel(
            model_module,
            device_ids=[local_rank],
            output_device=local_rank,
            find_unused_parameters=True,
        )
    else:
        if device.type == "cuda":
            if num_gpus is not None:
                if num_gpus < 1:
                    raise ValueError("num_gpus must be at least 1 or None")
            num_cuda_devices = torch.cuda.device_count()
            if num_cuda_devices < 1:
                raise Exception("CUDA devices are not available.")
            if gpu_ids is None:
                num_gpus = (
                    num_cuda_devices
                    if num_gpus is None
                    else min(num_gpus, num_cuda_devices)
                )
                gpu_ids = list(range(num_gpus))
            else:
                gpu_ids = list(
                    set(list(range(num_cuda_devices))).intersection(gpu_ids))
            if len(gpu_ids) > 0:
                model = torch.nn.DataParallel(model_module, device_ids=gpu_ids)
    return model