from functools import reduce
from itertools import combinations
from textwrap import dedent
from typing import Set, override
from unittest import TestCase

from puzzles._solver import Solver

Point = tuple[int, int]


class DayEightSolver(Solver):
    width: int
    height: int
    antenna_arrays: dict[str, set[Point]]

    @override
    def __init__(self, puzzle_input: str):
        self.antenna_arrays = {}

        lines = puzzle_input.strip().split("\n")
        self.width = len(lines[0])
        self.height = len(lines)

        for x, line in enumerate(puzzle_input.strip().split("\n")):
            for y, char in enumerate(line):
                if char.isalnum():
                    if char not in self.antenna_arrays:
                        self.antenna_arrays[char] = set()
                    self.antenna_arrays[char].add((x, y))

    @override
    def solve_p1(self) -> int:
        return len(
            reduce(
                Set.union,
                (
                    self.antinodes(a, b, harmonic=False)
                    for antenna_array in self.antenna_arrays.values()
                    for a, b in combinations(antenna_array, 2)
                ),
            )
        )

    @override
    def solve_p2(self) -> int:
        return len(
            reduce(
                Set.union,
                (
                    self.antinodes(a, b, harmonic=True)
                    for antenna_array in self.antenna_arrays.values()
                    for a, b in combinations(antenna_array, 2)
                ),
            )
        )

    def antinodes(self, a: Point, b: Point, harmonic=False) -> set[Point]:
        points = set()

        dx = b[0] - a[0]
        dy = b[1] - a[1]

        if harmonic:
            antinode = a
            while self.is_in_grid(antinode):
                points.add(antinode)
                antinode = (antinode[0] - dx, antinode[1] - dy)

            antinode = b
            while self.is_in_grid(antinode):
                points.add(antinode)
                antinode = (antinode[0] + dx, antinode[1] + dy)

        else:
            antinode = (a[0] - dx, a[1] - dy)
            if self.is_in_grid(antinode):
                points.add(antinode)

            antinode = (b[0] + dx, b[1] + dy)
            if self.is_in_grid(antinode):
                points.add(antinode)

        return points

    def is_in_grid(self, point: Point) -> bool:
        return 0 <= point[0] < self.height and 0 <= point[1] < self.width


class TestDayEightSolver(TestCase):
    def test(self):
        solver = DayEightSolver(
            dedent(
                """
                ............
                ........0...
                .....0......
                .......0....
                ....0.......
                ......A.....
                ............
                ............
                ........A...
                .........A..
                ............
                ............
                """
            )
        )
        self.assertEqual(solver.width, 12)
        self.assertEqual(solver.height, 12)
        self.assertSetEqual(
            solver.antenna_arrays["0"],
            {(1, 8), (2, 5), (3, 7), (4, 4)},
        )
        self.assertEqual(solver.solve_p1(), 14)
        self.assertEqual(solver.solve_p2(), 34)