import torch


def arrange_neighbor(neighbor1, neighbor2, assignment):
    """This function re-arrange the neighbor, so the neighbors will be matched

    Args:
        neighbor1: 3D tensor, [bs, num_node1, neighbor_k, dim]
        neighbor2: 3D tensor, [bs, num_node2, neighbor_k, dim]
        assignment: [bs, num_node1, num_node2, neighbor_k, 2]
    """
    index1, index2 = torch.chunk(assignment, 2, dim=-1) # [bs, num_node1, num_node2, neighbor_k, 1]
    index1 = index1.repeat(1, 1, 1, 1, neighbor1.size(-1)) # [bs, num_node1, num_node2, neighbor_k, dim]
    index2 = index2.repeat(1, 1, 1, 1, neighbor2.size(-1)) # [bs, num_node1, num_node2, neighbor_k, dim]
    neighbor1 = torch.gather(neighbor1, index=index1.long(), dim=-2)
    neighbor2 = torch.gather(neighbor2, index=index2.long(), dim=-2)

    return neighbor1, neighbor2