import torch


def calc_distances_to_means(means, tensors, distance_function='cosine'):
    """
    Function to calculate distances between tensors, in our case the mean zs per class and z for each input.
    Wrapper around torch.nn.functonal distances with specification of which distance function to choose.

    Parameters:
        means (list): List of length corresponding to number of classes containing torch tensors (typically mean zs).
        tensors (list): List of length corresponding to number of classes containing tensors (typically zs).
        distance_function (str): Specification of distance function. Choice of cosine|euclidean|mix.

    Returns:
        list: List of length corresponding to number of classes containing tensors with distance values
    """

    def distance_func(a, b, w_eucl, w_cos):
        """
        Weighted distance function consisting of cosine and euclidean components.

        Parameters:
            a (torch.Tensor): First tensor.
            b (torch.Tensor): Second tensor.
            w_eucl (float): Weight for the euclidean distance term.
            w_cos (float): Weight for the cosine similarity term.
        """
        d = w_cos * (1 - torch.nn.functional.cosine_similarity(a.view(1, -1), b)) + \
            w_eucl * torch.nn.functional.pairwise_distance(a.view(1, -1), b, p=2)
        return d

    distances = []

    # weight values for individual distance components
    w_eucl = 0.0
    w_cos = 0.0
    if distance_function == 'euclidean':
        w_eucl = 1.0
    elif distance_function == 'cosine':
        w_cos = 1.0
    elif distance_function == 'mix':
        w_eucl = 0.5
        w_cos = 0.5
    else:
        raise ValueError("distance function not implemented")

    # loop through each class in means and calculate the distances with the respective tensor.
    for i in range(len(means)):
        # check for tensor type, e.g. list could be empty
        if isinstance(tensors[i], torch.Tensor) and isinstance(means[i], torch.Tensor):
            distances.append(distance_func(means[i], tensors[i], w_eucl, w_cos))
        else:
            distances.append([])

    return distances