import json


def load_ontology(dataset, ontology_file=None):
    """
    Read ontology file for event to argument mapping.
    """
    ontology_dict = {}
    if not ontology_file:  # use the default file path
        if not dataset:
            raise ValueError
        with open("event_role_{}.json".format(dataset), "r") as f:
            ontology_dict = json.load(f)
    else:
        with open(ontology_file, "r") as f:
            ontology_dict = json.load(f)

    for evt_name, evt_dict in ontology_dict.items():
        for i, argname in enumerate(evt_dict["roles"]):
            evt_dict["arg{}".format(i + 1)] = argname
            # argname -> role is not a one-to-one mapping
            if argname in evt_dict:
                evt_dict[argname].append("arg{}".format(i + 1))
            else:
                evt_dict[argname] = ["arg{}".format(i + 1)]

    return ontology_dict