import torch
from typing import Tuple


def remove_pairs_from_same_domain(
    batch_domain_indices: torch.Tensor, 
    a_pair_info: torch.Tensor, 
    b_pair_info: torch.Tensor
    ) -> Tuple[torch.Tensor,torch.Tensor]:
    """
    In training, we want only pairs from different domains. We
    enforce that their feature embeddings are similar.

    We could have 1 million sampled pairs from a minibatch of size 5.
    (Number of elements in batch (batch_domain_indices) need not
    agree with number of sampled pairs!)

        Args:
        -   batch_domain_indices: Tensor of shape (K,), for each example
                in minibatch, which domain did it come from.
        -   a_pair_info:  (M,3) array representing (bi,hi,wi)
                where these represent (batch index, row index, column index)
                into a NCHW tensor for samples A.
        -   b_pair_info: (M,3) as above, but for samples B. (a,b) are paired

        Returns:
        -   a_pair_info: (N,3), where N <= M (discarded same domain elements)
        -   b_pair_info: (N,3), where N <= M
    """
    batch_dim_a_idxs = a_pair_info[:,0]
    batch_dim_b_idxs = b_pair_info[:,0]
    # remove locations with identical domains in pos/neg pairs
    a_domain = batch_domain_indices[batch_dim_a_idxs]
    b_domain = batch_domain_indices[batch_dim_b_idxs]

    is_valid_pair = (a_domain != b_domain).nonzero().squeeze()
    return a_pair_info[is_valid_pair], b_pair_info[is_valid_pair]