import numpy as np
from numpy.lib.stride_tricks import as_strided

from day11.common import char_to_int


def get_filters(arr):
    KERNEL_SHAPE = (3, 3)
    return as_strided(
        arr,
        shape=(
            arr.shape[0] - KERNEL_SHAPE[0] + 1,  # The feature map is a few pixels smaller than the input
            arr.shape[1] - KERNEL_SHAPE[1] + 1,
            KERNEL_SHAPE[0],
            KERNEL_SHAPE[1],
        ),
        strides=(
            arr.strides[0],
            arr.strides[1],
            arr.strides[0],  # When we move one step in the 3rd dimension, we should move one step in the original data too
            arr.strides[1],
        ),
        writeable=False,  # totally use this to avoid writing to memory in weird places
    )


def parse(arr):
    if arr[1, 1] == 0:
        return 0
    if np.count_nonzero(arr == 1) == 0:
        return 1
    if arr[1, 1] == 1 and np.count_nonzero(arr == 1) >= 5:
        return -1
    return arr[1, 1]


lines = np.array([[char_to_int(x) for x in line.strip()] for line in open("input.txt")], dtype=np.byte)
lines = np.pad(lines, 1, 'constant', constant_values=0)
shape = np.array(lines.shape)
last = -1
while True:
    curr = np.count_nonzero(lines == 1)
    if curr == last:
        print(last)
        break
    last = curr
    expanded_input = get_filters(lines)
    new_input = np.zeros(shape, dtype=np.byte)
    for idx in np.ndindex(*shape-2):
        target = tuple(np.array(idx) + 1)
        new_input[target] = parse(expanded_input[idx])
    lines = new_input