#!/usr/bin/env python3

BUG = '#'
EMPTY = '.'
LEVEL = '.'*12 + '?' + '.'*12

edges = {7  : range(0, 5),
         11 : range(0, 21, 5),
         13 : range(4, 25, 5),
         17 : range(20, 25)}


def neighbours(lvl, i):
    r = []
    if i > 4:
        r.append(i-5)
    if i < 20:
        r.append(i+5)
    if (i % 5) < 4:
        r.append(i+1)
    if (i % 5) > 0:
        r.append(i-1)
    return list(zip(4*[lvl], r))

def rec_neigh(lvl, i):
    ns = neighbours(lvl, i)
    for inner, outer in edges.items():
        if i == inner:
            ns += [(lvl+1, j) for j in outer]
        elif i in outer:
            ns += [(lvl-1, inner)]
    return ns

def fate(val, ns):
    if val == BUG and ns.count(BUG) != 1:
        return EMPTY
    elif val == EMPTY and ns.count(BUG) in (1, 2):
        return BUG
    return val

def trim(stack):
    while BUG not in stack[-1]:
        stack.pop()
    while BUG not in stack[0]:
        stack.pop(0)

def next(bug_stack, nfunc=rec_neigh):
    bug_stack = 2*[LEVEL] + bug_stack + 2*[LEVEL]
    new = []
    for lvl in range(1, len(bug_stack)-1):
        new.append('')
        for i, val in enumerate(bug_stack[lvl]):
            ns = nfunc(lvl, i)
            nvals = [bug_stack[l][j] for l, j in ns]
            new[-1] += fate(val, nvals)
    trim(new)
    return new

def bugs_sim(bug_stack, i):
    for i in range(i):
        bug_stack = next(bug_stack)
    return bug_stack

def biodiversity(bugs):
    return sum(2**pw for pw, c in enumerate(bugs) if c == BUG)

def preproc(puzzle_input):
    return puzzle_input.replace('\n', '')

def partI(bugs):
    layouts = set()
    b = bugs
    while b not in layouts:
        layouts.add(b)
        b = next([b], neighbours)[0]

    return biodiversity(b)

def partII(bugs):
    bugs = bugs[:12] + '?' + bugs[13:]
    stack = [bugs]
    return ''.join(bugs_sim(stack, 200)).count(BUG)