#!/usr/bin/env python3

# https://adventofcode.com/2024/day/8

from itertools import combinations


def get_lines(filename: str) -> list:
    with open(filename, "r") as file:
        return [line.strip() for line in file.readlines()]


def get_antenna_groups(lines: list) -> dict:
    antenna_groups = {}
    for row, line in enumerate(lines):
        for col, char in enumerate(line):
            if char != ".":
                if char not in antenna_groups:
                    antenna_groups[char] = [(row, col)]
                else:
                    antenna_groups[char].append((row, col))
    return antenna_groups


def get_anti_nodes(pair: tuple, max_row: int, max_col: int, resonant: bool) -> set:
    anti_nodes = set()
    a, b = pair
    a_row, a_col = a
    b_row, b_col = b
    if resonant:
        anti_nodes.add((a_row, a_col))
        anti_nodes.add((b_row, b_col))
    delta_row = b_row - a_row
    delta_col = b_col - a_col
    loop = True
    a_invalid = False
    b_invalid = False
    while loop:
        anti_a_row = a_row - delta_row
        anti_a_col = a_col - delta_col
        anti_b_row = b_row + delta_row
        anti_b_col = b_col + delta_col
        if 0 <= anti_a_row <= max_row and 0 <= anti_a_col <= max_col:
            anti_nodes.add((anti_a_row, anti_a_col))
        else:
            a_invalid = True
        if 0 <= anti_b_row <= max_row and 0 <= anti_b_col <= max_col:
            anti_nodes.add((anti_b_row, anti_b_col))
        else:
            b_invalid = True
        if a_invalid and b_invalid:
            loop = False
        if resonant:
            a_row, a_col = anti_a_row, anti_a_col
            b_row, b_col = anti_b_row, anti_b_col
        else:
            loop = False
    return anti_nodes


def get_all_anti_nodes(antenna_groups: dict, max_row: int, max_col: int, resonant: bool) -> set:
    anti_nodes = set()
    for antenna_group in antenna_groups:
        antennas = antenna_groups[antenna_group]
        antenna_pairs = combinations(antennas, 2)
        for pair in antenna_pairs:
            new_anti_nodes = get_anti_nodes(pair, max_row, max_col, resonant)
            anti_nodes.update(new_anti_nodes)
    return anti_nodes


def main():
    # lines = get_lines("sample-input.txt")
    lines = get_lines("input.txt")

    antenna_groups = get_antenna_groups(lines)
    max_row = len(lines) - 1
    max_col = len(lines[0]) - 1
    anti_nodes = get_all_anti_nodes(antenna_groups, max_row, max_col, False)
    print(f"Part 1: There are {len(anti_nodes)} anti-nodes on the map.")
    anti_nodes = get_all_anti_nodes(antenna_groups, max_row, max_col, True)
    print(f"Part 2: There are {len(anti_nodes)} resonant anti-nodes on the map.")


if __name__ == '__main__':
    main()