from tools.aoc import AOCDay
from typing import Any


def simulate_generation(state: str, transform: dict, shift: int) -> (str, int):
    while not state.startswith("....."):
        state = '.' + state
        shift += 1

    while not state.endswith('.....'):
        state += '.'

    new_state = '..'
    for x in range(len(state) - 5):
        key = state[x:x+5]
        if key not in transform:
            new_state += '.'
        else:
            new_state += transform[key]

    while new_state.startswith('.'):
        new_state = new_state[1:]
        shift -= 1

    return new_state, shift


def get_pot_value(state: str, shift: int) -> int:
    sum = 0
    for i, x in enumerate(state):
        if x == '#':
            sum += i - shift

    return sum


class Day(AOCDay):
    inputs = [
        [
            (325, "input12_test"),
            (3494, "input12"),
        ],
        [
            (2850000002454, "input12"),
        ]
    ]

    def parse_input(self) -> (str, dict):
        state = self.getInput()[0].split(": ")[1]

        transform = {}
        for line in self.getInput()[2:]:
            k, v = line.split(" => ")
            transform[k] = v

        return state, transform

    def part1(self) -> Any:
        state, transform = self.parse_input()
        shift = 0
        for x in range(20):
            state, shift = simulate_generation(state, transform, shift)

        return get_pot_value(state, shift)

    def part2(self) -> Any:
        cache = {}
        num_gens = 50_000_000_000
        state, transform = self.parse_input()
        cache[state] = (0, 0)
        shift = 0
        for x in range(1, num_gens + 1):
            state, shift = simulate_generation(state, transform, shift)
            if state in cache:
                return get_pot_value(state, (-num_gens + 11))
            cache[state] = (x, shift)


if __name__ == '__main__':
    day = Day(2018, 12)
    day.run(verbose=True)