from typing import Dict
from typing import Optional
from typing import Set


def _remove_zero_statistics(
    statistics: Dict[str, Dict[str, int]],
    force_keep: Optional[Set[str]] = None,
    require_trivial_children: bool = False,
) -> Dict[str, Dict[str, int]]:
    """
    Any module that has zero for all available statistics is removed from the
    set of statistics. This can help declutter the reporting of statistics
    if many submodules have zero statistics. Assumes the statistics have
    a model hierarchy starting with a root that has name ''.

    Args:
        statistics (dict(str, dict(str, int))) : the statistics to
            remove zeros from. Organized as a dictionary over modules,
            which are each a dictionary over statistic types.
        force_keep (set(str) or None) : a set of modules to always keep, even
            if they are all zero.
        require_trivial_children (bool) : If True, a statistic will only
            be deleted if all its children are also deleted. Defaults to
            False.

    Returns:
        dict(str, dict(str, int)) : the input statistics dictionary,
            with submodules removed if they have zero for all statistics.
    """
    out_stats: Dict[str, Dict[str, int]] = {}
    _force_keep: Set[str] = force_keep if force_keep else set() | {""}

    def keep_stat(name: str) -> None:
        prefix = name + ("." if name else "")
        trivial_children = True
        for mod in statistics:
            # 'if mod' excludes root = '', which is never a child
            if mod and mod.count(".") == prefix.count(".") and mod.startswith(prefix):
                keep_stat(mod)
                trivial_children &= mod not in out_stats

        if (
            (not all(val == 0 for val in statistics[name].values()))
            or (name in _force_keep)
            or (require_trivial_children and not trivial_children)
        ):
            out_stats[name] = statistics[name].copy()

    keep_stat("")
    return out_stats