from pathlib import Path


def find_files(base_root: Path):
    """
    Search the given base directory for the actual dataset root.
    This makes it a little easier for the dataset manager
    :param base_root:
    :return:
    """
    # These are folders we expect within the dataset folder structure. If we hit them, we've gone too far
    excluded_folders = {'depth', 'rgb', '__MACOSX'}
    to_search = {Path(base_root)}
    while len(to_search) > 0:
        candidate_root = to_search.pop()

        # Make the derivative paths we're looking for. All of these must exist.
        rgb_path = candidate_root / 'rgb.txt'
        trajectory_path = candidate_root / 'groundtruth.txt'
        depth_path = candidate_root / 'depth.txt'

        # If all the required files are present, return that root and the file paths.
        if rgb_path.is_file() and trajectory_path.is_file() and depth_path.is_file():
            return candidate_root, rgb_path, depth_path, trajectory_path

        # This was not the directory we were looking for, search the subdirectories
        for child_path in candidate_root.iterdir():
            if child_path.is_dir() and child_path.name not in excluded_folders:
                to_search.add(child_path)

    # Could not find the necessary files to import, raise an exception.
    raise FileNotFoundError("Could not find a valid root directory within '{0}'".format(base_root))