from itertools import count
from utils import open_day
import numpy as np

with open_day(11) as f:
    inp = [list(map(int, line.rstrip())) for line in f]

inp = np.array(inp)

def adjacentify(points):
    tot = np.zeros(points.shape, dtype=int)
    for dy in (-1, 0, 1):
        for dx in (-1, 0, 1):
            if dx == 0 and dy == 0: continue
            p = points.copy()
            if dy:
                p = np.roll(p, dy, 0)
                p[min(dy, 0), :] = 0
            if dx:
                p = np.roll(p, dx, 1)
                p[:, min(dx, 0)] = 0
            tot += p
    return tot

flashes = 0
for step in count():
    if step == 100: print(flashes)
    flashed = np.zeros(inp.shape, dtype=bool)
    inp += 1
    while True:
        flash = (inp > 9) & ~flashed
        if not flash.any(): break
        inp += adjacentify(flash * 1)
        flashed |= flash
    inp[flashed] = 0
    s = flashed.sum()
    flashes += s
    if s == flashed.shape[0] * flashed.shape[1]:
        print(step + 1)
        break