from aocd.models import Puzzle
from aocd import submit
from dataclasses import dataclass
import math
import re


def parse_input(input_data: str) -> dict:
    """
    Parse input and filter
    - first get positions of all numbers in field
    - read their value if there is an adjacent symbol
    Assumptions to greatly simplify the task (thx reddit):
    - No number touches more than one symbol
    - No non-star symbol touches exactly two numbers
    """

    # convert input into a 2d array for easier parsing
    field = [line for line in input_data.split("\n")]

    # create a dictionary of all symbols and their coordinates
    symbols = {
        (i, char): []
        for i in range(140)
        for char in range(140)
        if field[i][char] not in "01234566789."
    }

    # for each row in field
    for i, row in enumerate(field):
        # for each number in row create a dictionary of any number with coords
        for n in re.finditer(r"\d+", row):
            numbers = {
                (i, char)
                for i in (i - 1, i, i + 1)
                for char in range(n.start() - 1, n.end() + 1)
            }

            # append each number to the list of numbers
            for o in numbers & symbols.keys():
                symbols[o].append(int(n.group()))
    return symbols


if __name__ == "__main__":
    # get puzzle and parse data
    puzzle = Puzzle(year=2023, day=3)

    # parse_input(puzzle.input_data)
    nums = parse_input(puzzle.input_data)

    # part a:
    answer_a = sum(sum(p) for p in nums.values())
    print(f"{answer_a}")
    submit(answer_a, part="a", day=3, year=2023)

    # part b:
    answer_b = sum(math.prod(p) for p in nums.values() if len(p) == 2)
    print(f"{answer_b}")
    submit(answer_b, part="b", day=3, year=2023)