from typing import List
from typing import Dict
from typing import Callable


def bucket_molecules(molecules: List[Dict], keys: List[Callable]) -> Dict:
    """
    Organize molecules into nested dictionaries according to molecule properties
    specified by the "keys".

    The nested dictionary has keys based on the functions used in "keys",
    and the innermost value is a list. For example, if
    `keys = [
        lambda x: x["formula_alphabetical"],
        lambda x: len(x.get("bonds", list())),
        lambda x: x["charge"]
    ]`,
    then the returned bucket dictionary is something like:

    bucket[formula][num_bonds][charge] = [mol_entry1, mol_entry2, ...]

    where mol_entry1, mol_entry2, ... have the same formula, number of bonds, and charge.

    Args:
        entries (List[Dict]): a list of molecule entries to bucket
        keys (List[Callable]): each function should return a property of the molecule

    Returns:
        Dict of molecule entry grouped according to keys.
    """

    num_keys = len(keys)
    buckets = dict()  # type: ignore
    for m in molecules:
        b = buckets
        for i, j in enumerate(keys):
            v = j(m)
            if i == num_keys - 1:
                b.setdefault(v, []).append(m)
            else:
                b.setdefault(v, {})
            b = b[v]

    return buckets