import torch
from typing import Tuple


def find_matching_pairs(
    y: torch.Tensor,
    a_pair_info: torch.Tensor,
    b_pair_info: torch.Tensor) -> Tuple[torch.Tensor,torch.Tensor]:
    """
    Given a batch of ground truth label maps, and sampled pixel
    pair locations (pairs are across label maps), identify which 
    pairs are matching vs. non-matching and return corresponding metadata
    (basically, partition them).

        Args:
        -   y: Tensor of size (B,H,W) representing 2-d label maps
                for B examples.
        -   a_pair_info:
        -   b_pair_info:

        Returns:
        -   pos_pair_info: Pytorch tensor containing info about each positive pair (a,b). Contains
                (a batch_idx, a row, a col, b batch_idx, b row, b col)
        -   neg_pair_info: Same as above, but for negative pairs.
    """
    batch_dim_a_idxs = a_pair_info[:,0]
    px_1d_a_y = a_pair_info[:,1]
    px_1d_a_x = a_pair_info[:,2]

    batch_dim_b_idxs = b_pair_info[:,0]
    px_1d_b_y = b_pair_info[:,1]
    px_1d_b_x = b_pair_info[:,2]

    # extract category indices
    cls_vals_a = y[batch_dim_a_idxs, px_1d_a_y, px_1d_a_x]
    cls_vals_b = y[batch_dim_b_idxs, px_1d_b_y, px_1d_b_x]

    # compare category indices for equality
    is_same_class = (cls_vals_a == cls_vals_b).nonzero().squeeze()
    is_diff_class = (cls_vals_a != cls_vals_b).nonzero().squeeze()

    a_pos_info = a_pair_info[is_same_class]
    a_neg_info = a_pair_info[is_diff_class]

    b_pos_info = b_pair_info[is_same_class]
    b_neg_info = b_pair_info[is_diff_class]

    pos_pair_info = torch.cat([a_pos_info, b_pos_info], dim=1)
    neg_pair_info = torch.cat([a_neg_info, b_neg_info], dim=1)

    return pos_pair_info, neg_pair_info