import matrix
import shared
import itertools
import functools
from pprint import pprint
from collections import defaultdict

SYMBOLS="*%@#+-/$=&"

def get_all_numbers_and_starting_coords(mat):
    """
        for a coordinate, walks right until it encounters a non digit, 
        when it does, it adds that number string to a collection, 
        and starts over, walking right until it continues to next row
    """
    nums = []
    for y, row in enumerate(mat):
        x = 0
        while x < len(row):
            num = matrix.number_starting_at(mat, x,y)
            if num is None:
                x += 1
                continue
            nums.append(((x,y), num))
            x += len(num)
    return nums

def get_row_coords(x1, y, length):
    """
    for a given x,y point, it just creates the rest of the x/y coords to the right of this coord for a run
    """
    coords = []
    for x in range(x1, x1+length):
        coords.append((x,y))
    return coords

def process_number(mx, coords, num):
    line_coords = get_row_coords(*coords, len(num))
    for coord in line_coords:
        any_symbols = [(cs, v) for cs,v in matrix.get_neighbor_coords(mx, *coord) if v in SYMBOLS]
        if any_symbols:
            return True, any_symbols[0]
    return False, None

# @shared.profile
def part1(mat):
    choose = []
    nums = get_all_numbers_and_starting_coords(mat)
    for coords, num in nums:
        valid, symbol_coord = process_number(mat, coords, num)
        if valid:
            choose.append(num)
    total = 0
    for n in choose:
        total += int(n)
    print(total)


# @shared.profile
def part2(mat):
    splats = matrix.find_in_matrix(mat, "*", False)
    _nums = get_all_numbers_and_starting_coords(mat)

    symbols = defaultdict(list)
    nums = get_all_numbers_and_starting_coords(mat)
    for coords, num in nums:
        valid, sc = process_number(mat, coords, num)
        if valid:
            symbols[(sc[0]['r'],sc[0]['c'])].append(num)
    
    total = 0
    for splat in splats:
        neighbors = symbols[splat]
        if len(neighbors) != 2:
            continue
        ratio = int(neighbors[0]) * int(neighbors[1])
        total += ratio
    print(total)



def main():
    mat = shared.load_file_char_matrix(shared.get_fname(3))
    with shared.elapsed_timer() as elapsed:
        part1(mat)
        print("🕒", elapsed())

    with shared.elapsed_timer() as elapsed:
        part2(mat)
        print("🕒", elapsed())


if __name__ == "__main__":
    main()