import numpy as np

def read_map(filename: str):
	# Read file as list of strings
	with open(filename, "r") as fp:
		data = fp.read().splitlines()

	# Use Numpy to make it a 2D array of characters
	tp_str_map = np.array(data, dtype=str)\
		.view("U1")\
		.reshape((len(data), -1))

	# Make it a boolean map
	tp_map = np.zeros_like(tp_str_map, dtype=bool)
	tp_map[tp_str_map == "@"] = True
	return tp_map

def print_map(tp_map):
	view = np.zeros_like(tp_map, dtype="U1")
	view[tp_map] = "@"
	view[~tp_map] = "."
	for i in range(view.shape[0]):
		for j in range(view.shape[1]):
			print(view[i,j], end="")
		print()

def count_neighbours(tp_map):
	delta1D = np.array([-1, 0, 1])
	X, Y = np.meshgrid(delta1D, delta1D)
	neighbour_counts = np.zeros_like(tp_map, dtype=int)

	for x, y in np.ndindex(tp_map.shape):
		if not tp_map[x,y]:
			continue
		dX, dY = X.copy() + x, Y.copy() + y
		legal_idx = (dX >= 0)*(dX < X.shape[0])*(dY >= 0)*(dY < Y.shape[1])
		neighbour_counts[x, y] = (tp_map[dX[legal_idx], dY[legal_idx]]).sum()
	print(neighbour_counts)
	max_rolls = 4
	print(neighbour_counts < max_rolls)
	return (neighbour_counts < max_rolls).sum()

if __name__ == "__main__":
	test_map = read_map("testinput")
	print(test_map)
	print(count_neighbours(test_map))