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,
) -> 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.

    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.

    Returns:
        dict(str, dict(str, int)) : the input statistics dictionary,
            with submodules removed if they have zero for all statistics.
    """
    out_stats = {}
    if force_keep is None:
        force_keep = set()
    for mod, stats in statistics.items():
        if not all(val == 0 for val in stats.values()) or mod in force_keep:
            out_stats[mod] = stats.copy()
    return out_stats