import torch


def generate_final_context_matrix(context_matrix_dict, ending_key_to_exclude="_nokg"):
    """Takes the average of the context matrices where their key does not end
    in ending_key_to_exclude.

    Args:
        context_matrix_dict: Dict of batch x M x K x H matrices of embeddings for each candidate
        ending_key_to_exclude: key to exclude from average

    Returns:
        tensor average of call tensors in context_matrix_dict unless key ends in ending_key_to_exclude
    """
    new_ctx = []
    for key, val in context_matrix_dict.items():
        if not key.endswith(ending_key_to_exclude):
            new_ctx.append(val)
    assert len(new_ctx) > 0, (
        f"You have provided a context matrix dict with only keys ending with _nokg. We average the context matrices "
        f"that do not end in _nokg as the final context matrix. Please rename the final matrix."
    )
    return torch.sum(torch.stack(new_ctx), dim=0) / len(new_ctx)