from collections import defaultdict


def gen_counts_old(input_lines, interaction_types=None, residuelabels=None):
    """
    Parse each line in `input_lines` as a line from get_*_contacts and return interaction-counts for each residue pair.
    If `residuelabels` is defined it is used to modify residue identifiers and to filter out residues not indicated.

    For example:
        inputs = [
            "# total_frames: 3",
            "\t".join(["0", "hbbb", "A:ALA:1:N", "A:ARG:4:O"]),
            "\t".join(["0", "vdw", "A:ALA:1:CB", "A:ARG:4:CA"]),
            "\t".join(["1", "vdw", "A:ALA:1:N", "A:CYS:5:CA"]),
            "\t".join(["2", "hbbb", "A:THR:2:N", "A:CYS:5:O"]),
            "\t".join(["2", "hbss", "A:ALA:1:N", "A:CYS:5:O"])
        ]
        labels = {"A:ALA:1": "A1", "A:ARG:4": "R4", "A:CYS:5": "C5"}

        # Only consider hbbb and vdw, filter away THR, and map to single-letter labels
        gen_counts(inputs, ["hbbb", "vdw"], labels)
        #  => { ("A1", "R4"): 1, ("A1", "C5"): 1 }

    Parameters
    ----------
    input_lines: Iterable[str]
        Interactions formatted as get_*_contacts output, e.g. ["0\thbbb\tA:ALA:1:N\tA:ARG:4:H", ...]
    interaction_types: list of str
        Which interaction types to consider
    residuelabels: dict of (str: str)
        Remaps and filters residuelabels, e.g. {"A:ARG:4": "R4"}

    Returns
    -------
    (int, dict of (str, str): int)
        Total frame-count and mapping of residue-residue interactions to frame-count
    """

    def atomid_to_resid(atom):
        return atom[0:atom.rfind(":")]
        # return ':'.join(atom.split(':')[1:3])

    # Maps residue pairs to set of frames in which they're present
    rescontact_frames = defaultdict(set)
    total_frames = 0

    for line in input_lines:
        line = line.strip()
        if "total_frames" in line:
            tokens = line.split(" ")
            total_frames = int(tokens[1][tokens[1].find(":")+1:])

        if len(line) == 0 or line[0] == "#":
            continue

        tokens = line.split("\t")

        # Check that the interaction type is specified
        if interaction_types and tokens[1] not in interaction_types:
            continue

        frame = int(tokens[0])
        if frame + 1 > total_frames:
            total_frames = frame + 1

        res1 = atomid_to_resid(tokens[2])
        res2 = atomid_to_resid(tokens[3])

        # Change residue id according to `residuelabels` or skip if any of the residues are not present
        if residuelabels is not None:
            if res1 not in residuelabels or res2 not in residuelabels:
                continue
            res1 = residuelabels[res1]
            res2 = residuelabels[res2]

        # Ensure lexicographical order of residue names
        if res2 < res1:
            res1, res2 = res2, res1

        rescontact_frames[(res1, res2)].add(frame)

    # Insted of returning list of frames for each interaction, only return number of frames
    rescontact_counts = {(res1, res2): len(frames) for (res1, res2), frames in rescontact_frames.items()}
    return total_frames, rescontact_counts