from typing import Union
from typing import List
import yaml
import re


def load_yaml(
    filename: str, keys_to_list: bool = True
) -> Union[dict, List[dict]]:
    """Load in YAML file.

    Args:
        filename (str): YAML filename to load from.
        keys_to_list (bool): Option to transform dict of eval entries into list.

    Returns:
        List[dict]: List of evaluation results.
    """
    loader = yaml.SafeLoader
    loader.add_implicit_resolver(
        "tag:yaml.org,2002:float",
        re.compile(
            """^(?:
        [-+]?(?:[0-9][0-9_]*)\\.[0-9_]*(?:[eE][-+]?[0-9]+)?
        |[-+]?(?:[0-9][0-9_]*)(?:[eE][-+]?[0-9]+)
        |\\.[0-9_]+(?:[eE][-+][0-9]+)?
        |[-+]?[0-9][0-9_]*(?::[0-5]?[0-9])+\\.[0-9_]*
        |[-+]?\\.(?:inf|Inf|INF)
        |\\.(?:nan|NaN|NAN))$""",
            re.X,
        ),
        list("-+0123456789."),
    )
    with open(filename) as file:
        yaml_temp = yaml.load(file, Loader=loader)
    if keys_to_list:
        # From dict of evals to list of evals
        yaml_log = []
        for k in yaml_temp.keys():
            yaml_log.append(yaml_temp[k])
        return yaml_log
    else:
        return yaml_temp