import aocd
import math
import collections
import itertools

def solve(data):
    grid = collections.defaultdict(list)
    antinodes = set()

    rows = data.split('\n')
    width = len(rows[0])
    height = len(rows)

    for y, row in enumerate(rows):
        for x, cell in enumerate(row):
            if cell == '.':
                continue

            for prev_x, prev_y in grid[cell]:
                dx = x - prev_x
                dy = y - prev_y

                angle = math.gcd(dx, dy)
                dx, dy = dx // angle, dy // angle

                # In the direction of the previous antenna...
                for i in itertools.count(step=-1):
                    if 0 <= (new_x := prev_x + dx*i) < width and 0 <= (new_y := prev_y + dy*i) < height:
                        antinodes.add((new_x, new_y))
                    else:
                        break

                # In the direction of the current antenna...
                for i in itertools.count(step=1):
                    if 0 <= (new_x := x + dx*i) < width and 0 <= (new_y := y + dy*i) < height:
                        antinodes.add((new_x, new_y))
                    else:
                        break
            
            grid[cell].append((x, y))
    
    return sum(1 for _ in antinodes)

if __name__ == '__main__':
    puz = aocd.models.Puzzle(year=2024, day=8)

    # data = puz.examples[0].input_data
    data = puz.input_data

    solution = solve(data)

    print(solution)
    aocd.submit(solution, year=2024, day=8, part='b')