from scipy.signal import convolve2d
from scipy.ndimage.filters import convolve1d
input_serial_num = 1718

def my_convolve2d(matrix, size=3):
    kernel = np.ones(size)
    intermediate = convolve1d(matrix, kernel, axis=0,
                              mode='constant',origin=(size-1)//2)[:matrix.shape[1]-(size-1):]
    return convolve1d(intermediate, kernel, axis=1,
                     mode='constant', origin=(size-1)//2)[:,:matrix.shape[0]-(size-1)]

def hundreds_digit(x):
    return x//100 - x//100//10*10

hdv = np.vectorize(hundreds_digit)

def power_matrix(serial_num, shape=(300,300)):
    rackid = (np.arange(shape[0])+1)+10
    y_vals = np.arange(shape[1])+1
    return hdv((rackid[:,np.newaxis]*y_vals[np.newaxis,:] + serial_num)*rackid[:,np.newaxis])-5


def find_largest_square(matrix, shape=(3,3)):
    #kernel = np.ones(shape)
    #convolved_matrix = convolve2d(matrix, kernel, mode='valid')
    convolved_matrix = my_convolve2d(matrix, shape[0])
    indicies = np.unravel_index(np.argmax(convolved_matrix), convolved_matrix.shape)
    power = convolved_matrix[indicies]
    return np.array(np.unravel_index(np.argmax(convolved_matrix), convolved_matrix.shape))+1, power

find_largest_square(power_matrix(input_serial_num))

#------Part 2------
def find_largest_square_shape(matrix):
    max_power = 0
    max_indicies = None
    max_size = 0
    for square_shape in range(1,301):
        #print(square_shape)
        indicies, power = find_largest_square(matrix, shape=(square_shape, square_shape))
        max_indicies = indicies if power>max_power else max_indicies
        max_size = square_shape if power>max_power else max_size
        max_power = power if power>max_power else max_power
    return max_indicies, max_size, max_power

find_largest_square_shape(power_matrix(input_serial_num))