import torch


def one_hot(targets, num_classes, dtype=torch.float32, device=None):
    """
    Encode the targets (an tensor of integers representing a class)
    as one hot encoding.

    Support target as N-dimensional data (e.g., 3D segmentation map).

    Equivalent to torch.nn.functional.one_hot for backward compatibility with pytorch 1.0

    Args:
        num_classes: the total number of classes
        targets: a N-dimensional integral tensor (e.g., 1D for classification, 2D for 2D segmentation map...)
        dtype: the type of the output tensor
        device: the device of the one-hot encoded tensor. If `None`, use the target's device

    Returns:
        a one hot encoding of a N-dimentional integral tensor
    """
    if device is None:
        device = targets.device

    nb_samples = len(targets)
    if len(targets.shape) == 2:
        # 2D target (e.g., classification)
        encoded_shape = (nb_samples, num_classes)
    else:
        # N-d target (e.g., segmentation map)
        encoded_shape = tuple([nb_samples, num_classes] + list(targets.shape[1:]))

    with torch.no_grad():
        encoded_target = torch.zeros(encoded_shape, dtype=dtype, device=device)
        encoded_target.scatter_(1, targets.unsqueeze(1), 1)
    return encoded_target