from collections import Counter
from copy import copy
from typing import List, Tuple, Dict, FrozenSet

from aoc import BaseAssignment

class UnknownException(Exception):
    pass

Item = FrozenSet[str]
Information = Tuple[List[Item], List[Item]]

class Assignment(BaseAssignment):
    def parse_item(self, item: str) -> Tuple[List[Item], List[Item]]:
        input, output = item.split('|')

        return (
            [frozenset(i.strip()) for i in input.split(' ') if i != ''],
            [frozenset(i.strip()) for i in output.split(' ') if i != '']
        )

    def read_input(self, example = False) -> List[Tuple[List[str], List[str]]]:
        return list(super().read_input(example))

    @staticmethod
    def calc_digit(item: Item, known_numbers: Dict[int, Item] = None) -> int:
        if known_numbers is None:
            known_numbers = {}

        match len(item):
            case 2:
                return 1
            case 3:
                return 7
            case 4:
                return 4
            case 5:
                if 1 in known_numbers and item.issuperset(known_numbers[1]):
                    return 3
                elif 4 in known_numbers and len(item | known_numbers[4]) == 7:
                    return 2
                else:
                    return 5
            case 6:
                if 1 in known_numbers and not item.issuperset(known_numbers[1]):
                    return 6
                elif 3 in known_numbers and item.issuperset(known_numbers[3]):
                    return 9
                else:
                    return 0
            case 7:
                return 8
            case _:
                raise UnknownException()

    @staticmethod
    def unique_filter(item: Item) -> bool:
        try:
            return Assignment.calc_digit(item) in [1, 4, 7, 8]
        except UnknownException:
            return False


class AssignmentOne(Assignment):
    example_result = 26

    def run(self, input: List[Information]) -> int:
        return sum([
            len(list(filter(Assignment.unique_filter, output)))
            for _, output in input
        ])


class AssignmentTwo(Assignment):
    example_result = 61229

    @staticmethod
    def find_all_numbers(items: List[Item]) -> Dict[int, Item]:
        known_numbers = {
            Assignment.calc_digit(item): item
            for item in filter(Assignment.unique_filter, items)
        }

        unknown_numbers = {
            item
            for item
            in filter(
                lambda item: item not in known_numbers.values(),
                items
            )
        }

        while len(unknown_numbers) > 0:
            for number in copy(unknown_numbers):
                try:
                    nr = Assignment.calc_digit(number, known_numbers)
                    known_numbers[nr] = number
                    unknown_numbers.remove(number)
                except UnknownException:
                    continue

        return known_numbers

    def run(self, input: List[Information]) -> int:
        return sum([
            sum(
                self.calc_digit(item, self.find_all_numbers(input + output))
                    * (10 ** (len(output) - index - 1))
                for index, item in enumerate(output)
            )
            for input, output in input
        ])