import torch


def euclidean_dist(x, y):
    """
    Calculates Euclidean distance (square of it)
    :param x: size [n_query_total, out_dim=1600] - queries
    :param y: size [n_ways, out_dim=1600] - prototypes
    """
    n = x.size(0)  # total number of query points = n_query_total
    m = y.size(0)  # number of classes = n_ways
    d = x.size(1)  # dimension of pic embedding = 1600 for mini-ImageNet
    if d != y.size(1):
        raise ValueError(f'Pic embedding for prototype {y.size(1)} and query {d} data arent equal')

    x = x.unsqueeze(1).expand(n, m, d)  # size = [n_query_total, n_ways, 1600]
    y = y.unsqueeze(0).expand(n, m, d)  # size = [n_query_total, n_ways, 1600]

    return torch.pow(x - y, 2).sum(2)