# -*- coding: utf-8 -*-
from abc import ABC
from functools import reduce
from typing import Optional, Iterator

from aoc import BaseAssignment, I, T
from aoc.datastructures import Coordinate


class Assignment(BaseAssignment, ABC):
    def symbol_overlap(
        self, number_coordinates: list[Coordinate], symbols: set[Coordinate]
    ) -> set[Coordinate]:
        neighbours = {
            neighbour
            for coordinate in number_coordinates
            for neighbour in coordinate.neighbours()
        }

        return neighbours & symbols

    def read_input(
        self, example=False
    ) -> tuple[list[str], list[tuple[list[Coordinate], int]], set[Coordinate]]:
        schematic = list(super().read_input(example))
        numbers: list[tuple[list[Coordinate], int]] = list()
        symbols: set[Coordinate] = set()

        for y, row in enumerate(schematic):
            number: Optional[list[list[Coordinate], str]] = None

            for x, item in enumerate(row):
                if item.isdigit():
                    if number is None:
                        number = [[], ""]

                    number[0].append(Coordinate(x, y))
                    number[1] += item
                    continue

                if number is not None:
                    numbers.append((number[0], int(number[1])))
                    number = None

                if item != ".":
                    symbols.add(Coordinate(x, y))

            if number is not None:
                numbers.append((number[0], int(number[1])))

        return schematic, numbers, symbols


class AssignmentOne(Assignment):
    example_result = 4361

    def is_part_number(
        self, number_coordinates: list[Coordinate], symbols: set[Coordinate]
    ) -> bool:
        return len(self.symbol_overlap(number_coordinates, symbols)) > 0

    def run(
        self, input: tuple[list[tuple[list[Coordinate], int]], set[Coordinate]]
    ) -> int:
        _, numbers, symbols = input

        part_numbers = []

        for coordinates, number in numbers:
            if self.is_part_number(coordinates, symbols):
                part_numbers.append(number)

        return sum(part_numbers)


class AssignmentTwo(Assignment):
    example_result = 467835

    def run(
        self, input: tuple[list[tuple[list[Coordinate], int]], set[Coordinate]]
    ) -> int:
        schematic, numbers, symbols = input

        possible_gears = {}

        for coordinates, number in numbers:
            for symbol in self.symbol_overlap(coordinates, symbols):
                if schematic[symbol.y][symbol.x] != "*":
                    continue

                if symbol not in possible_gears:
                    possible_gears[symbol] = []

                possible_gears[symbol].append(number)

        return sum(
            [
                reduce(lambda total, number: total * number, numbers, 1)
                for symbol, numbers in possible_gears.items()
                if len(numbers) > 1
            ]
        )