from itertools import combinations
from tools.aoc import AOCDay
from typing import Any
from tools.grid import Grid


def get_antinode_count(grid: Grid, repeat: bool = False) -> int:
    antinodes = set()
    frequencies = set(grid.get(x) for x in grid.getActiveCells())
    for frequency in frequencies:
        antennas = set(x for x in grid.getActiveCells() if grid.get(x) == frequency)
        if len(antennas) < 2:
            continue
        if repeat:
            antinodes |= antennas

        for pair in combinations(antennas, 2):
            a, b = (pair[0], pair[1]) if pair[0] < pair[1] else (pair[1], pair[0])
            diff = b - a

            while grid.isWithinBoundaries(a - diff):
                antinodes.add(a - diff)
                if not repeat:
                    break
                a -= diff

            while grid.isWithinBoundaries(b + diff):
                antinodes.add(b + diff)
                if not repeat:
                    break
                b += diff

    return len(antinodes)


class Day(AOCDay):
    inputs = [
        [
            (14, "input8_test"),
            (244, "input8_dennis"),
            (413, "input8"),
        ],
        [
            (34, "input8_test"),
            (912, "input8_dennis"),
            (1417, "input8"),
        ],
    ]

    def parse_input(self) -> Grid:
        grid = Grid.from_data(self.getInput(), default=False, translate={".": False})
        grid.minX, grid.minY, grid.maxX, grid.maxY = 0, 0, len(self.getInput()[0]) - 1, len(self.getInput()) - 1
        return grid

    def part1(self) -> Any:
        grid = self.parse_input()
        return get_antinode_count(grid, repeat=False)

    def part2(self) -> Any:
        grid = self.parse_input()
        return get_antinode_count(grid, repeat=True)


if __name__ == "__main__":
    day = Day(2024, 8)
    day.run(verbose=True)