import tqdm


def generate_formula_dict(materials_store, query=None):
    """
    Function that generates a nested dictionary of structures
    keyed first by formula and then by task_id using
    mongo aggregation pipelines

    Args:
        materials_store (Store): store of materials

    Returns:
        Nested dictionary keyed by formula-mp_id with structure values.

    """
    props = ["pretty_formula", "structure", "task_id", "magnetic_type"]
    results = list(materials_store.groupby("pretty_formula", properties=props,
                                           criteria=query))
    formula_dict = {}
    for result in tqdm.tqdm(results):
        formula = result['_id']['pretty_formula']
        task_ids = [d['task_id'] for d in result['docs']]
        structures = [d['structure'] for d in result['docs']]
        formula_dict[formula] = dict(zip(task_ids, structures))
    return formula_dict