import numpy as np
from day11.common import char_to_int

seats = np.array([[char_to_int(x) for x in line.strip()] for line in open("input.txt")], dtype=np.byte)
neighbors = np.zeros_like(seats, dtype=np.byte)
next_neighbors = np.zeros_like(seats, dtype=np.byte)
RAYS = [(1, 0), (-1, 1), (0, 1), (1, 1)]
occupied = -1


def find_neighbor(idx, ray):
    idx = tuple(map(sum, zip(idx, ray)))
    while 0 <= idx[0] < seats.shape[0] and 0 <= idx[1] < seats.shape[1]:
        if seats[idx] == -1:
            return 0, idx
        elif seats[idx] == 0:
            idx = tuple(map(sum, zip(idx, ray)))
            continue
        else:
            return 1, idx
    else:
        return None


while True:
    new_occupied = np.count_nonzero(seats == 1)
    if occupied == new_occupied:
        break
    occupied = new_occupied
    result = np.zeros_like(seats)
    for idx, val in np.ndenumerate(seats):
        if val == 0:
            continue
        findings = (find_neighbor(idx, ray) for ray in RAYS)
        findings = [f for f in findings if f is not None]
        neighbors[idx] += sum(occ for occ, f in findings)
        if val == -1 and neighbors[idx] == 0:
            result[idx] = 1
        elif val == 1 and neighbors[idx] >= 5:
            result[idx] = -1
        else:
            result[idx] = seats[idx]
        if result[idx] == 1:
            for occ, f in findings:
                next_neighbors[f] += 1
    seats = result
    neighbors = next_neighbors
    next_neighbors = np.zeros_like(seats, dtype=np.byte)
    print(occupied)

print(occupied)