import math
from abc import ABC
from collections import Counter
from copy import copy
from typing import List, Tuple, Union, Dict

from aoc import BaseAssignment


class Assignment(BaseAssignment, ABC):
    max_pre_calculations = 4
    steps = 0
    mapping_cache = dict()

    def parse_item(self, item: str) -> Union[str, Tuple[str, str], None]:
        if '>' in item:
            return item.split(' -> ')

        if item == '':
            return None

        return item


    def read_input(self, example = False) -> Tuple[str, Dict[str, str]]:
        input = super().read_input(example)

        template = next(input)
        next(input)

        mapping = dict(input)

        return template, mapping

    def run(self, input: Tuple[str, Dict[str, str]]) -> int:
        template, mapping = input

        counter = Counter([
            f'{template[i:i+2]}'
            for i in range(len(template) - 1)
        ])

        char_count = {
            **{
                  c: 1 if c in template else 0
                  for c
                  in mapping.values()
            },
            **Counter(template),
        }

        for _ in range(self.steps):
            new_counter = {
                pair: 0
                for pair
                in mapping.keys()
            }

            for key, value in counter.items():
                middle = mapping[key]

                char_count[middle] += value
                new_counter[f'{key[0]}{middle}'] += value
                new_counter[f'{middle}{key[1]}'] += value

            counter = new_counter

        return max(char_count.values()) - min(char_count.values())

class AssignmentOne(Assignment):
    example_result = 1588
    steps = 10

class AssignmentTwo(Assignment):
    max_pre_calculations = 4
    example_result = 2188189693529
    steps = 40