#!/usr/bin/env python
# 2021 - Advent Of Code - 14 part 2

from collections import Counter


def parse_file(file):
    poly_list = []
    insertion_rules = {}
    with open(file, encoding="utf-8") as f:
        poly_str = list(f.readline().strip())

        for i, p in enumerate(poly_str):
            if i + 1 == len(poly_str):
                break
            pair = '' + p + poly_str[i + 1]
            poly_list.append(pair)
        f.readline()

        for line in f.readlines():
            pair, insert = line.strip().split(' -> ')
            insertion_rules[pair] = [pair[0] + insert, insert + pair[1]]

        f_l = poly_str[0], poly_str[-1]

        return poly_list, insertion_rules, f_l


def step_poly(init_count, ruleset):
    new_count = dict.fromkeys(ruleset.keys(), 0)
    for p in init_count.keys():
        nb = init_count[p]
        p1, p2 = ruleset[p]
        new_count[p1] += nb
        new_count[p2] += nb

    return new_count


# poly, rules, first_last = parse_file('input_example.txt')
poly, rules, first_last = parse_file('input.txt')
print(f'poly: {poly}')
print(f'rules: {rules}')

poly_count = dict.fromkeys(rules.keys(), 0)
for polymer in poly:
    poly_count[polymer] += 1
print(f'init count: {poly_count}')

for s in range(40):
    # print(f'step: {s}')
    poly_count = step_poly(poly_count, rules)

print(f'poly count: {poly_count}')

elements = set([k for key in poly_count.keys() for k in (key[0], key[1])])
print(f'elements: {elements}')
elem_count = dict.fromkeys(elements, 0)

for polymer in poly_count.keys():
    elem_count[polymer[0]] += poly_count[polymer]
    elem_count[polymer[1]] += poly_count[polymer]

# remove double count, and fix first and last element
elem_count[first_last[0]] += 1
elem_count[first_last[1]] += 1
elem_count = {k: int(v/2) for k, v in elem_count.items()}
print(f'fixed elem count: {elem_count}')

c = Counter(elem_count).most_common()[0][1] - Counter(elem_count).most_common()[-1][1]
print(c)