import torch


def calc_gt_indices(batch_keys_gt, batch_assignments_gt):
    """
        calc_gt_indices() calculate the ground truth indices and number of valid key points.
        Arguments:
            batch_keys_gt: [tensor B x N x 3] the last column stores the indicator of the key points
                            if it is larger than 0 it is valid, otherwise invalid
            batch_assignments_gt [tensor B x N x N], the ground truth assignment matrix
        Returns:
            indices_gt: [tensor B x N ]: from source to target, each element stores the index of target matches
            key_num_gt: [tensor B]: number of valid key points per input of the batch
    """
    _, indices_gt = torch.max(
        batch_assignments_gt, 2
    )  # get ground truth matches from source to target
    indices_gt += (
        1
    )  # remember that indices start counting from 1 for 0 is used to store empty key points
    mask_gt = (batch_keys_gt[:, :, 2] > 0).long()  # get the valid key point masks
    indices_gt = indices_gt * mask_gt
    key_num_gt = mask_gt.sum(dim=1).float()
    return indices_gt, key_num_gt