from collections import defaultdict


def gen_frequencies(count_list):
    """
    Take a list of residue contact counts (see output of `gen_counts`) and compute total counts and frequencies.

    Example:
        clist = [
            (4, {("A1", "R4"): 4, ("A1", "C5"): 3}),  # First simulation has 4 frames and two contacts
            (3, {("A1", "R4"): 2})                    # Second simulation has 3 frames and one contact
        ]
        gen_frequencies(clist)
        # Returns: (7, {("A1", "R4"): (6, 0.857), ("A1", "C5"): (3, 0.429)})

    Parameters
    ----------
    count_list: list of (int, dict of (str, str): int)
        List with individual frame counts and dictionaries mapping residue pairs to frame-counts

    Return
    ------
    (int, dict of (str, str): (int, float))
        Total framecount and mapping of residue ID pairs to the number of frames in which they contact and the frequency
    """
    # print("gen_frequencies")
    # print(count_list)
    rescontact_count = defaultdict(int)
    total_frames = 0
    for frames, rescount_dict in count_list:
        total_frames += frames

        for (res1, res2), count in rescount_dict.items():
            rescontact_count[(res1, res2)] += count

    respair_freqs = {respair: (count, float(count) / total_frames) for respair, count in rescontact_count.items()}
    return total_frames, respair_freqs