import torch


def calc_mto(batch_assignments, src_indices_gt, src_key_num_gt):
    """
    calc_mto() calculate the one-to-many matching score, notice one is source, many is destination
    Arguments:
        batch_assignment [tensor float B x 32 x 32]: the batch of the predicted assignment matrix
        src_indices_gt [tensor long B x 32 ]: the batch of ground truth key point indices from source to target
        src_key_num_gt [tensor float Bx 1]: the ground truth number of valid key points for a batch with batch size B
    Returns:
        mto [tensor B x 1] cpu: the mto score for the batch    
    """
    values, indices = torch.max(batch_assignments, 2)
    indices += (
        1
    )  # remember that indices start counting from 1 for 0 is used to store empty key points
    mask = (src_indices_gt == indices).long()
    indices *= mask
    num_unique = torch.tensor([float(len(torch.unique(kk))) - 1 for kk in indices])
    mto = 1 - torch.div(num_unique, src_key_num_gt.cpu())
    return mto