from itertools import combinations
from typing import NamedTuple

test_input = """
............
........0...
.....0......
.......0....
....0.......
......A.....
............
............
........A...
.........A..
............
............
"""

test_solution_p1 = 14
test_solution_p2 = 34


def solve_p1(puzzle_input: str) -> int:
    grid = _parse_signal_grid(puzzle_input)
    antinodes = _calculate_antinodes(grid)
    # _print_grid(grid, antinodes)
    return len(antinodes)


def solve_p2(puzzle_input: str) -> int:
    grid = _parse_signal_grid(puzzle_input)
    antinodes = _calculate_antinodes(grid, harmonic=True)
    # _print_grid(grid, antinodes)
    return len(antinodes)


class Point(NamedTuple):
    x: int
    y: int


class SignalGrid(NamedTuple):
    width: int
    height: int
    antenna_arrays: dict[str, set[Point]]


def _parse_signal_grid(puzzle_input: str) -> SignalGrid:
    lines = puzzle_input.strip().split("\n")
    antenna_arrays: dict[str, set[Point]] = {}

    for x, line in enumerate(lines):
        for y, char in enumerate(line):
            if char.isalnum():
                if char not in antenna_arrays:
                    antenna_arrays[char] = set()
                antenna_arrays[char].add(Point(x, y))

    return SignalGrid(len(lines[0]), len(lines), antenna_arrays)


def _calculate_antinodes(grid: SignalGrid, harmonic=False) -> set[Point]:
    antinodes: set[Point] = set()

    for antenna_array in grid.antenna_arrays.values():
        for a, b in combinations(antenna_array, 2):
            dx = b.x - a.x
            dy = b.y - a.y

            # Left antinodes
            if harmonic:
                antinode = a
                while True:
                    antinodes.add(antinode)
                    antinode = Point(antinode.x - dx, antinode.y - dy)
                    if not _is_in_grid(grid, antinode):
                        break
            else:
                antinode = Point(a.x - dx, a.y - dy)
                if _is_in_grid(grid, antinode):
                    antinodes.add(antinode)

            # Right antinodes
            if harmonic:
                antinode = b
                while True:
                    antinodes.add(antinode)
                    antinode = Point(antinode.x + dx, antinode.y + dy)
                    if not _is_in_grid(grid, antinode):
                        break
            else:
                antinode = Point(b.x + dx, b.y + dy)
                if _is_in_grid(grid, antinode):
                    antinodes.add(antinode)

    return antinodes


def _is_in_grid(grid: SignalGrid, point: Point) -> bool:
    return 0 <= point.x < grid.height and 0 <= point.y < grid.width


def _print_grid(grid: SignalGrid, antinodes: set[Point]):
    output = [["."] * grid.width for _ in range(grid.height)]
    for antinode in antinodes:
        output[antinode.x][antinode.y] = "#"
    for frequency, antenna_array in grid.antenna_arrays.items():
        for antenna in antenna_array:
            output[antenna.x][antenna.y] = frequency
    print("\n".join("".join(row) for row in output))