#!/bin/python3
import sys,time,re
from pprint import pprint
sys.path.insert(0, '../../')
from fred import list2int,get_re,nprint,lprint,grid_valid,loadFile,subTuples,addTuples,toGrid,get_value_in_direction
start_time = time.time()

input_f = 'input'

part = 2
#########################################
#                                       #
#              Part 1                   #
#                                       #
#########################################

if part == 1:
    def findAntinodes(grid:list,pos:set,current_node:set):
        antinodes = []
        for row in range(pos[0],len(grid)):
            for col in range(0,len(grid[1])):

                if get_value_in_direction(grid,(row,col)) == current_node and (row,col) != pos:

                    dist = subTuples((row,col),pos)
                    add = addTuples(dist,(row,col))
                    sub = subTuples(pos,dist)
                    if grid_valid(sub[0],sub[1],grid):
                        antinodes.append(sub)
                    if grid_valid(add[0],add[1],grid):
                        antinodes.append(add)
        
        return antinodes

    grid = toGrid(input_f)

    size = (len(grid),len(grid[0]))
    antinodes = []
    for row in range(0,size[0]):
        for col in range(0,size[1]):
            current_node = get_value_in_direction(grid,(row,col))
            if current_node != '.':
                antinodes += findAntinodes(grid,(row,col),current_node)
    print(len(set(antinodes)))


#########################################
#                                       #
#              Part 2                   #
#                                       #
#########################################
if part == 2:
    def findAntinodes(grid:list,pos:set,current_node:set):
        antinodes = []
        for row in range(0,len(grid)):
            for col in range(0,len(grid[1])):

                if get_value_in_direction(grid,(row,col)) == current_node and (row,col) != pos:

                    dist = subTuples((row,col),pos)
                    add = pos
                    sub = pos
                    while grid_valid(add[0],add[1],grid):
                        add = addTuples(dist,add)
                        if grid_valid(add[0],add[1],grid):
                            antinodes.append(add)
                    
                    while grid_valid(sub[0],sub[1],grid):
                        sub = subTuples(sub,dist)
                        if grid_valid(sub[0],sub[1],grid):
                            antinodes.append(sub)
        
        return antinodes
    
    grid = toGrid(input_f)

    size = (len(grid),len(grid[0]))
    antinodes = []
    for row in range(0,size[0]):
        for col in range(0,size[1]):
            current_node = get_value_in_direction(grid,(row,col))
            if current_node != '.':
                antinodes += findAntinodes(grid,(row,col),current_node)
    print(len(set(antinodes)))

print("--- %s seconds ---" % (time.time() - start_time))