import torch


def select_seeds(dist1: torch.Tensor, R1: float, scores1: torch.Tensor, fnn12: torch.Tensor, mnn: torch.Tensor):
    """
        Select seed correspondences among the set of available matches.

        dist1: Precomputed distance matrix between keypoints in image I_1
        R1: Base radius of neighborhoods in image I_1
        scores1: Confidence scores on the putative_matches. Usually holds Lowe's ratio scores.
        fnn12: Matches between keypoints of I_1 and I_2.
               The i-th entry of fnn12 is j if and only if keypoint k_i in image I_1 is matched to keypoint k_j in image I_2
        mnn: A mask indicating which putative matches are also mutual nearest neighbors. See documentation on 'force_seed_mnn' in the DEFAULT_CONFIG.
             If None, it disables the mutual nearest neighbor filtering on seed point selection.
             Expected a bool tensor with shape (num_keypoints_in_source_image,)

        Returns:
            Indices of seed points.

            im1seeds: Keypoint index of chosen seeds in image I_1
            im2seeds: Keypoint index of chosen seeds in image I_2
    """
    im1neighmap = dist1 < R1**2  # (n1, n1)
    # find out who scores higher than whom
    im1scorescomp = scores1.unsqueeze(1) > scores1.unsqueeze(0)  # (n1, n1)
    # find out who scores higher than all of its neighbors: seed points
    if mnn is not None:
        im1bs = (~torch.any(im1neighmap & im1scorescomp & mnn.unsqueeze(0),
                            dim=1)) & mnn & (scores1 < 0.8**2)  # (n1,)
    else:
        im1bs = (~torch.any(im1neighmap & im1scorescomp, dim=1)) & (scores1 <
                                                                    0.8**2)

    # collect all seeds in both images and the 1NN of the seeds of the other image
    im1seeds = torch.where(im1bs)[0]  # (n1bs) index format
    im2seeds = fnn12[im1bs]  # (n1bs) index format
    return im1seeds, im2seeds