import re


def to_int(bitlist):
    out = 0
    for bit in bitlist:
        out = out * 2 + bit
    return out


def grow(init, num_gens, progress=False):
    first, last = min(init), max(init)
    current_gen = set(init)
    for g in range(num_gens):
        next_gen = set()
        for i in range(first - 2, last + 3):
            key = to_int(1 if b in current_gen else 0 for b in range(i - 2, i + 3))
            if key in positives:
                next_gen.add(i)
                if i < first:
                    first = i
                if i > last:
                    last = i
        current_gen = next_gen

        if progress:
            print('Generations:', g + 1, end='\r')

    if progress: print('Generations:', g + 1)
    return current_gen


def test(n):
    pots = grow(initial, n)
    first, last = min(pots), max(pots)
    chars = []
    for i in range(first - 1, last + 2):
        chars.append('#' if i in pots else '.')
    print(''.join(chars))


with open('data/12.txt') as f:
    initial = re.search('[#.]+', next(f)).group(0)
    initial = {i for i, p in enumerate(initial) if p == '#'}
    next(f)
    
    positives, negatives = set(), set()
    for line in f:
        state, result = line.strip().split(' => ')
        s = to_int(int(c == '#') for c in state)
        if result == '#':
            positives.add(s)
        else:
            negatives.add(s)


pots = grow(initial, 20)
print(f'Part 1: {sum(pots)}\n')


# find pattern and exploit
first_stable = 0
prev_gen = []
for i in range(1, 1000):
    current_gen = sorted(grow(initial, i))
    if current_gen == [x + 1 for x in prev_gen]:
        first_stable = i - 1
        break
    else:
        prev_gen = current_gen
magic_number = sum(prev_gen)


def cheat(num_gens):
    if num_gens < first_stable:
        raise RuntimeError("Can't cheat.")
    return magic_number + (num_gens - first_stable) * len(prev_gen)

print(f'Part 2: {cheat(5 * 10 ** 10)}\n')