from typing import List
from typing import Dict
from typing import Set
import collections
import json


def load_entities_fn(file_name: str, types: List[str]) -> Dict[str, Set[str]]:
    """Given a json-lines file from Pubtator, loads the entities.

    Args:
        file_name: Path to the json-lines files from Pubtator.
        types: List of entity types to be loaded.

    Returns:
        type_entities_dict: A dictionary where key is an entity type and 
            value is a set of entities of that type.
    """
    type_entities_dict = collections.defaultdict(set)
    with open(file_name) as f:
        for line in f:
            json_obj = json.loads(line)
            for entity in json_obj['entity_list']:
                if entity['type'].lower() in types:
                    for entity_i in entity['name'].lower().split('|'):
                        type_entities_dict[entity['type'].lower()].add(
                            entity_i.strip())
    return type_entities_dict