import os
import pickle


def get_dataset_path(lmdb_base_dir: str,
                     dataset_name: str) -> str:
    """
    Get the path to the dataset in the lmdb directory
    Args:
        lmdb_base_dir: base directory of the lmdb dataset
        dataset_name: name of the dataset

    Returns:
        path to the dataset
    """

    dataset_path = os.path.join(lmdb_base_dir, dataset_name)
    if not os.path.exists(dataset_path):
        raise FileNotFoundError(f"Dataset {dataset_name} not found in {lmdb_base_dir}")
    return dataset_path


def load_key_stats(lmdb_dir: str) -> dict:
    """
    Load key statistics from the LMDB database.
    :param lmdb_dir: Directory of the LMDB database.
    :return: Dictionary with statistics.
    """

    key_stats_path = f"{lmdb_dir}/key_stats.pickle"

    if not os.path.exists(key_stats_path):
        return dict()

    with open(key_stats_path, "rb") as f:
        key_stats = pickle.load(f)
    return key_stats