import numpy as np
from part1 import read_map, print_map, count_neighbours, count_extractable

def remove_iteratively(tp_map):
	max_rolls = 4
	rolls_removed_counts = []
	rolls_removed_count = 1_000_000
	while rolls_removed_count > 0:
		neighbour_counts = count_neighbours(tp_map)
		rolls_to_remove = (neighbour_counts < max_rolls)*tp_map
		rolls_removed_count = rolls_to_remove.sum()
		rolls_removed_counts.append(rolls_removed_count)
		tp_map[rolls_to_remove] = False
	return sum(rolls_removed_counts)

if __name__ == "__main__":
	test_map = read_map("testinput")
	# print(test_map)
	assert count_extractable(test_map) == 13
	assert remove_iteratively(test_map.copy()) == 43
	# TODO: Compare the output maps.

	real_map = read_map("input")
	print(remove_iteratively(real_map.copy()))