from tools.aoc import AOCDay
from typing import Any


PRIMES = [2, 3, 5, 7, 11, 13, 17]


def get_max_power_point(power_grid: list[list[int]], window_size: int = 3) -> (int, int, int):
    max_power_level = 0
    ret = (0, 0, 0)

    for y in range(300 - window_size):
        for x in range(300 - window_size):
            sub_power_level = sum(
                sum(power_grid[z][x:x+window_size])
                for z in range(y, y + window_size)
            )
            if sub_power_level > max_power_level:
                max_power_level = sub_power_level
                ret = (x, y, max_power_level)

    return ret


class Day(AOCDay):
    inputs = [
        [
            ("33,45", "input11_test"),
            ("21,61", "input11_test2"),
            ("243,34", "input11"),
        ],
        [
            ("90,269,16", "input11_test"),
            ("232,251,12", "input11_test2"),
            ("90,214,15", "input11"),
        ]
    ]

    def get_power_grid(self) -> list[list[int]]:
        serial = self.getInput(return_type=int)
        grid = []
        for y in range(300):
            sub_grid = []
            for x in range(300):
                rack_id = x + 10
                power_level = rack_id * y
                power_level += serial
                power_level *= rack_id
                power_level = (power_level // 100) % 10
                power_level -= 5
                sub_grid.append(power_level)
            grid.append(sub_grid)

        return grid

    def part1(self) -> Any:
        power_grid = self.get_power_grid()
        return "%s,%s" % get_max_power_point(power_grid)[:-1]

    def part2(self) -> Any:
        power_grid = self.get_power_grid()
        max_power_level = 0
        max_power_ret = (0, 0, 0)
        for i in range(2, 301):
            (x, y, max_power) = get_max_power_point(power_grid, i)
            if (x, y, max_power) == (0, 0, 0):
                break
            if max_power > max_power_level:
                max_power_level = max_power
                max_power_ret = (x, y, i)

        return "%s,%s,%s" % max_power_ret


if __name__ == '__main__':
    day = Day(2018, 11)
    day.run(verbose=True)