import torch


def extract_local_patterns(
        fnn12: torch.Tensor,
        fnn_to_seed_local_consistency_map_corr: torch.Tensor, k1: torch.Tensor,
        k2: torch.Tensor, im1seeds: torch.Tensor, im2seeds: torch.Tensor,
        scores: torch.Tensor):
    """
        Prepare local neighborhoods around each seed for the parallel RANSACs. This involves two steps:
            1) Collect all selected keypoints and refer them with respect to their seed point
            2) Sort keypoints by score for the progressive sampling to pick the best samples first

        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
        fnn_to_seed_local_consistency_map_corr: Boolean matrix of size (num_seeds, num_keypoints).
                                                Entry (i, j) is True iff keypoint j was assigned to seed i.
        k1: Keypoint locations in image I_1
        k2: Keypoint locations in image I_2
        im1seeds: Keypoint index of chosen seeds in image I_1
        im2seeds: Keypoint index of chosen seeds in image I_2
        scores: Scores to rank correspondences by confidence.
                Lower scores are assumed to be more confident, consistently with Lowe's ratio scores.
                Note: scores should be between 0 and 1 for this function to work as expected.

        Returns:
            All information required for running the parallel RANSACs.
            Data is formatted so that all inputs for different RANSACs are concatenated
                along the same dimension to support different input sizes.

            im1loc: Keypoint locations in image I_1 for each RANSAC sample.
            im2loc: Keypoint locations in image I_2 for each RANSAC sample.
            ransidx: Integer identifier of the RANSAC problem.
                     This allows to distinguish inputs belonging to the same problem.
            tokp1: Index of the original keypoint in image I_1 for each RANSAC sample.
            tokp2: Index of the original keypoint in image I_2 for each RANSAC sample.
    """
    # first get an indexing representation of the assignments:
    # - ransidx holds the index of the seed for each assignment
    # - tokp1 holds the index of the keypoint in image I_1 for each assignment 
    ransidx, tokp1 = torch.where(fnn_to_seed_local_consistency_map_corr)
    # - and of course tokp2 holds the index of the corresponding keypoint in image I_2
    tokp2 = fnn12[tokp1]

    # Now take the locations in the image of each considered keypoint ... 
    im1abspattern = k1[tokp1]
    im2abspattern = k2[tokp2]

    # ... and subtract the location of its corresponding seed to get relative coordinates
    im1loc = im1abspattern - k1[im1seeds[ransidx]]
    im2loc = im2abspattern - k2[im2seeds[ransidx]]

    # Finally we need to sort keypoints by scores in a way that assignments to the same seed are close together
    # To achieve this we assume scores lie in (0, 1) and add the integer index of the corresponding seed
    expanded_local_scores = scores[tokp1] + ransidx.type(scores.dtype)

    sorting_perm = torch.argsort(expanded_local_scores)

    im1loc = im1loc[sorting_perm]
    im2loc = im2loc[sorting_perm]
    tokp1 = tokp1[sorting_perm]
    tokp2 = tokp2[sorting_perm]

    return im1loc, im2loc, ransidx, tokp1, tokp2