from typing import Union
import torch
from typing import List
import random


def get_noise(batch_size: int, latent_dimension, p_mixed_noise: float = 0.9, device: str = 'cuda') -> Union[
    torch.Tensor, List]:
    """
    Function returns an input noise for the style gan 2 generator.
    Iter a list of two noise vectors or one noise vector will be returned.
    :param batch_size: (int) Batch size to be used
    :param latent_dimension: (int) Latent dimensions to be utilized
    :param p_mixed_noise: (int) Probability that a mixed noise will be returned
    :param device: (str) Device to be utilized
    :return: List of noise tensors or single noise tensor
    """
    if (p_mixed_noise > 0) and (random.random() < p_mixed_noise):
        return list(torch.randn(2, batch_size, latent_dimension, dtype=torch.float32, device=device).unbind(0))
    else:
        return torch.randn(batch_size, latent_dimension, dtype=torch.float32, device=device)