#!/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) -> set:
    anti_nodes = set()
    a, b = pair
    a_row, a_col = a
    b_row, b_col = b
    vektor_row = b_row - a_row
    vektor_col = b_col - a_col
    anti_a_row = a_row - vektor_row
    anti_a_col = a_col - vektor_col
    if 0 <= anti_a_row <= max_row and 0 <= anti_a_col <= max_col:
        anti_nodes.add((anti_a_row, anti_a_col))
    anti_b_row = b_row + vektor_row
    anti_b_col = b_col + vektor_col
    if 0 <= anti_b_row <= max_row and 0 <= anti_b_col <= max_col:
        anti_nodes.add((anti_b_row, anti_b_col))
    return anti_nodes


def print_field(antennas: list, new_anti_nodes: set, max_row: int, max_col: int) -> None:
    field = []
    for i in range(max_row + 1):
        row = []
        for j in range(max_col + 1):
            row.append(".")
        field.append(row)
    for antenna in antennas:
        row, col = antenna
        field[row][col] = "A"
    for anti_node in new_anti_nodes:
        row, col = anti_node
        if field[row][col] == "A":
            field[row][col] = "@"
        else:
            field[row][col] = "#"
    for row in field:
        for char in row:
            print(char, end="")
        print()
    print()


def get_all_anti_nodes(antenna_groups: dict, max_row: int, max_col: int) -> 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)
            anti_nodes.update(new_anti_nodes)
        # print_field(antennas, anti_nodes, max_row, max_col)
    return anti_nodes


def main():
    # lines = get_lines("sample-input.txt")
    # lines = get_lines("sample-input2.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)
    print(f"Part 1: There are {len(anti_nodes)} anti-nodes on the map.")


if __name__ == '__main__':
    main()