import aocd
import collections

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

                for potential in ((prev_x-dx, prev_y-dy), (x+dx, y+dy)):
                    if 0 <= potential[0] < width and 0 <= potential[1] < height:
                        antinodes.add(potential)
            
            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='a')