from sys import argv
from scipy.ndimage import uniform_filter
from scipy.signal import fftconvolve
import numpy

try:
    SERIAL = int(argv[1])
except (ValueError, IndexError):
    print(f"Usage: {argv[0]} [serial]")

# Build the grid.
grid = numpy.zeros((300, 300), dtype=numpy.int32)
for y, row in enumerate(grid):
    rack_id = numpy.arange(1, 301, dtype=numpy.int32) + 10
    power_level = (rack_id * (y + 1) + SERIAL) * rack_id
    power_level = power_level // 100 % 10 - 5
    grid[y] = power_level

def main(from_, to=None):
    max_y, max_x = None, None
    max_size = 0
    max_value = 0

    if to is None:
        to = from_ + 1

    for size in range(from_, to):
        # Calculate the totals by convolving with a filter with ones.
        filter_ = numpy.ones((size, size), dtype=int)
        totals = fftconvolve(grid, filter_, "valid")

        # Find the indices of the maximum value.
        y, x = numpy.unravel_index(totals.argmax(), totals.shape)
        value = totals[y, x]

        if value > max_value:
            max_size, max_y, max_x = size, y, x
            max_value = value

    if to == from_ + 1:
        # If the range is one, it was part 1.
        print(f"{max_x + 1},{max_y + 1}")
    else:
        print(f"{max_x + 1},{max_y + 1},{max_size}")


print("part 1:")
main(3)
print("part 2:")
main(1, 100)