from typing import List
from typing import Dict
import torch


def average_checkpoints(
        all_state_dicts: List[Dict],
        out_path: str,
        ckpt_template: dict,
):
    """
    Average the state_dicts of top k checkpoints.

    Parameters
    ----------
    all_state_dicts
        A list of Pytorch state_dicts.
    out_path
        The path to save the averaged checkpoint.
    ckpt_template
        A dictionary of checkpoint template used during the training.

    Returns
    -------
    The averaged state_dict.
    """
    avg_state_dict = dict()
    for key in all_state_dicts[0]:
        arr = [state_dict[key] for state_dict in all_state_dicts]
        avg_state_dict[key] = sum(arr) / len(arr)

    ckpt_template["state_dict"] = avg_state_dict
    torch.save(ckpt_template, out_path)

    return avg_state_dict