import torch
from typing import List
import random


def get_indices(targets: torch.Tensor, per_class: bool = True,
                count: int = 4, seed: int = 0) -> List[int]:
    """Samples indices to visualize

    :param targets: ground truths
    :type targets: torch.Tensor
    :param per_class: whether count means count per class, defaults to True
    :type per_class: bool
    :param count: number of audios to return, defaults to 4
    :type count: int, optional
    :param seed: random seed, defaults to 0
    :type seed: int, optional
    """
    if not per_class:
        indices = list(range(len(targets)))
        indices = random.Random(seed).sample(indices, count)
    else:
        indices = []
        unique_targets, inverse_indices = torch.unique(
            targets, return_inverse=True)
        for index in range(len(unique_targets)):
            target_indices = torch.nonzero(
                inverse_indices == index, as_tuple=True)
            target_indices = target_indices[0].tolist()
            _indices = random.Random(seed).sample(
                target_indices, min(count, len(target_indices)))
            indices.extend(_indices)

    return indices