import collections
import torch


def get_filters(examples, n_relations):
    """
    Create filtering lists for evaluation.

    :param examples: Numpy array of size n_examples x 3 containing KG triples
    :param n_relations: Int indicating the total number of relations in the KG
    :return: lhs_final: Dictionary mapping queries (entity, relation) to filtered entities for
    left-hand-side prediction
    :return: rhs_final: Dictionary mapping queries (entity, relation) to filtered entities for
    right-hand-side prediction
    """
    lhs_filters = collections.defaultdict(set)
    rhs_filters = collections.defaultdict(set)
    for lhs, rel, rhs in examples:
        lhs_filters[(lhs, rel)].add(rhs)
        rhs_filters[(rhs, rel + n_relations)].add(lhs)
    lhs_final, rhs_final = {}, {}
    for k, v in lhs_filters.items():
        lhs_final[k] = torch.LongTensor(sorted(list(v)))
    for k, v in rhs_filters.items():
        rhs_final[k] = torch.LongTensor(sorted(list(v)))
    return lhs_final, rhs_final