import numpy as np

def flash(grid, flashed, idx_row, idx_col):
    current = grid[idx_row][idx_col]
    flashes = 0
    min_row = idx_row - 1
    max_row = idx_row + 2 #adding 2 because python range upper bound is not counted
    min_col = idx_col - 1
    max_col = idx_col + 2

    if min_row < 0:
        min_row = 0
    if max_row > len(grid):
        max_row = len(grid)

    if min_col < 0:
        min_col = 0
    if max_col >= len(grid[0]):
        max_col = len(grid[0])
    
    for i_row in range(min_row, max_row):
        for i_col in range(min_col, max_col):
            grid[i_row][i_col] += 1

def increment_all_by_one(grid):
    for idx_row, row in enumerate(grid):
        for idx_col, _ in enumerate(row):
            grid[idx_row][idx_col] += 1


def one_step(grid, flashed):
    flashes = 0
    while True:
        keep_checking_for_flash = False
        for idx_row, row in enumerate(grid):
            for idx_col, _ in enumerate(row):
                octopus = grid[idx_row][idx_col]
                if octopus > 9 and flashed[idx_row][idx_col] == False:
                    flashed[idx_row][idx_col] = True
                    flash(grid, flashed, idx_row, idx_col)
                    flashes += 1
                    grid[idx_row][idx_col] = -100
                    keep_checking_for_flash = True

        if not keep_checking_for_flash:
            break
    return flashes

def clean_up(grid, flashed):
    for idx_row, row in enumerate(grid):
        for idx_col, _ in enumerate(row):
            if flashed[idx_row][idx_col]:
                grid[idx_row][idx_col] = 0
                flashed[idx_row][idx_col] = False

def part1(lines):
    grid = [[int(num) for num in line] for line in lines]
    flashed = [[False for _ in range(10)] for _ in range(10)]
    flashes = 0

    for loop in range(100):
        increment_all_by_one(grid)
        flashes += one_step(grid, flashed)
        clean_up(grid, flashed)

    return flashes

def part2(lines):
    grid = [[int(num) for num in line] for line in lines]
    flashed = [[False for _ in range(10)] for _ in range(10)]
    step = 0

    while True:
        step += 1
        increment_all_by_one(grid)
        one_step(grid, flashed)
        clean_up(grid, flashed)
        if (np.array(grid)==0).all():
            break

    return step

if __name__ == "__main__":
    with open("../input.txt") as file:
        lines = [line.strip() for line in file]

    print(part1(lines))
    print(part2(lines))