""" Pair insertion """
import pprint
import collections

insertion_rules = {}

polymer_pairs = collections.defaultdict(int)

letter_count = collections.defaultdict(int)

with open("14_input.txt") as f:
    polymer_template = f.readline().strip()
    f.readline()

    for line in f:
        pair, insert = line.strip().split(" -> ")
        insertion_rules[pair] = insert

for i in range(len(polymer_template) - 1):
    pair = str(polymer_template[i:i+2])
    polymer_pairs[str(polymer_template[i:i+2])] += 1

print(polymer_pairs)

for _ in range(40):
    new_polymer_pairs = collections.defaultdict(int)
    for pair, n in polymer_pairs.items():
        insert = insertion_rules[pair]

        new_polymer_pairs[f"{pair[0]}{insert}"] += n
        new_polymer_pairs[f"{insert}{pair[0]}"] += n
        letter_count[insert] += n

    polymer_pairs = new_polymer_pairs

print(letter_count)
print(max(letter_count.values()) - min(letter_count.values()))