import torch


def get_prototypes(support_embeddings, support_labels, ways, shots):
    """
    Get prototypes of the support data in Prototypical Network. Given the support embeddings and labels, returns
    the prototypes of the relative classes in order.

    Args:
        support_embeddings(torch.Tensor): support embeddings calculated by embedding model. shape:
            (batch_size, embedding_size), batch_size should be ways * shots, embedding_size can be multiple dimensions.
        support_labels(torch.Tensor): labels of the embeddings. shape: (batch_size, 1). Support embedding number of
            each label should be the same and equal to shots.
        ways(int): ways of the task
        shots(int): shots of the task

    Returns:
        prototypes(torch.Tensor): shape: (ways, embedding_size). prototype[0] represents the prototype of label 0.
            According to Prototypical Network, the prototype of a label is the mean value of the support embeddings
            of the label.

    Examples:
        ..code-block:: python

            import torch
            from model import protonet

            support_embeddings = torch.Tensor([[1.1, 1.1, 1.1],
                                                   [0.0, 0.0, 0.0],
                                                   [0.9, 0.9, 0.9],
                                                   [0.0, 0.0, 0.0]])
            support_labels = torch.Tensor([[1], [0], [1], [0]])
            prototypes = protonet.get_prototypes(support_embeddings, support_labels, ways=2, shots=2)
            print(prototypes)  # Tensor(shape=[2, 3], ... [[0., 0., 0.], [1., 1., 1.]])

    """
    assert support_embeddings.shape[0] == ways * shots

    # support_embeddings[0].shape = embedding size
    prototypes = [torch.zeros_like(support_embeddings[0]) for _ in range(ways)]

    for i in range(len(support_embeddings)):  # range(support set 里的样本个数)
        proto_idx = int(support_labels.cpu().numpy()[i])
        prototypes[proto_idx] += support_embeddings[i]
    prototypes = [prototype / shots for prototype in prototypes]
    prototypes = torch.stack(prototypes)
    return prototypes