import torch


def _get_relation_types(dataset,):
    """
    Classify relations into 1-N, M-1, 1-1, M-N

    Bordes, Antoine, et al.
    "Translating embeddings for modeling multi-relational data."
    Advances in neural information processing systems. 2013.

    :return: dictionary mapping from int -> {1-N, M-1, 1-1, M-N}
    """
    relation_stats = torch.zeros((dataset.num_relations(), 6))
    for index, p in [
        (dataset.index("train_sp_to_o"), 1),
        (dataset.index("train_po_to_s"), 0),
    ]:
        for prefix, labels in index.items():
            relation_stats[prefix[p], 0 + p * 2] = labels.float().sum()
            relation_stats[prefix[p], 1 + p * 2] = (
                relation_stats[prefix[p], 1 + p * 2] + 1.0
            )
    relation_stats[:, 4] = (relation_stats[:, 0] / relation_stats[:, 1]) > 1.5
    relation_stats[:, 5] = (relation_stats[:, 2] / relation_stats[:, 3]) > 1.5
    result = dict()
    for i, relation in enumerate(dataset.relation_ids()):
        result[i] = "{}-{}".format(
            "1" if relation_stats[i, 4].item() == 0 else "M",
            "1" if relation_stats[i, 5].item() == 0 else "N",
        )
    return result