#!/usr/bin/python

import numpy as np



def init_grid(lines):
    dim = len(lines[0])
    grid = np.zeros((dim + 2, dim + 2), dtype="bool")

    for r, line in enumerate(lines):
        for c, char in enumerate(line):
            if char == "#":
                grid[r + 1, c + 1] = True

    return grid



def solve_partn(partn, init, steps):
    grid = init_grid(init)

    # Part 2: some lights are stuck in the "on" state
    if partn == 2:
        grid[([1, 1, -2, -2], [1, -2, 1, -2])] = True

    for i in range(steps):
        new_grid = grid.copy()
        for x in range(1, grid.shape[0] - 1):
            for y in range(1, grid.shape[1] - 1):
                neigh = np.sum(grid[x - 1 : x + 2, y - 1 : y + 2]) - grid[x, y]
                if grid[x, y]:
                    if neigh < 2 or neigh > 3:
                        new_grid[x, y] = False
                else: # not grid[x, y]
                    if neigh == 3:
                        new_grid[x, y] = True

        if partn == 2:
            new_grid[([1, 1, -2, -2], [1, -2, 1, -2])] = True

        grid = new_grid

    return np.count_nonzero(grid)



def main():
    # Read initial state from input text file
    with open("input.txt", "r") as f:
        lines = f.read().splitlines()

    print("Part 1 solution:", solve_partn(1, lines, 100)) # 821 for me
    print("Part 2 solution:", solve_partn(2, lines, 100)) # 886 for me



if __name__ == "__main__":
    main()