import collections


def clean_counts(probs):
    """
    Removes counts for certain SMARTS
    SMARTS are pairs of atoms
    Used only to prepare input for get_rxn_smarts

    Args:
        probs: OrderedDict of {SMARTS: probability}
               probability is actually a count of occurrences
               SMARTS are permutations of pairs of atoms (joined by bonds) where one of these atoms is not in a ring
               OrderedDict([('[#6]-[#6;!R]', 448), ('[#6]-[#7;!R]', 173)...])

    Returns:
        tot: sum of SMARTS counts
        probs2: an OrderedDict of {SMART: counts}

    """
    #            Triple bondedN, Carbonyl, Fl, Cl, Br, I
    exceptions = ["[#7]#", "[#8]=", "[#9]", "[#17]", "[#35]", "[#53]"]
    probs2 = collections.OrderedDict()

    # for key in probs:
    #   skip = False
    #   for exception in exceptions:
    #     if exception in key:
    #       tokens = re.split('\[|\]|;',key)
    #       alt_key = '['+tokens[3]+']'+tokens[2]+'['+tokens[1]+';!R]'
    #       probs[alt_key] += probs[key]

    for key in probs:
        skip = False
        for exception in exceptions:
            if exception in key:
                skip = True
        if not skip:
            probs2[key] = probs[key]

    tot = 0
    for key in probs2:
        tot += probs2[key]

    return tot, probs2