#!/usr/bin/env python

import numpy as np
from scipy.ndimage import label

with open("day9.txt") as data:
    depthmap = []
    for line in data:
        row = [int(num) for num in line.strip()]
        depthmap.append(row)
    depthmap = np.array(depthmap)

# part 1
# convolve across grid

def conv_filter(position: (int, int)) -> int:
    def neighbors():
        x, y = position
        neighboring = {(x-1, y), (x+1, y), (x, y-1), (x, y+1)}
        return {(xPos, yPos) for xPos, yPos in neighboring if 0 <= xPos <= depthmap.shape[0]-1 and 0 <= yPos <= depthmap.shape[1]-1}

    if all(depthmap[position] < depthmap[pos] for pos in neighbors()):
        return depthmap[position] + 1
    else:
        return 0

print(sum(conv_filter((x, y)) for x in range(depthmap.shape[0]) for y in range(depthmap.shape[1])))

# part 2
basins, num_basins = label(depthmap != 9)
counts = [np.count_nonzero(basins == num) for num in range(1, num_basins + 1)]
counts.sort()
print(np.prod(counts[-3:]))