from tools.aoc import AOCDay
from typing import Any


def getDisplayValue(input_str: str, output_str: str) -> int:
    inputs = input_str.split()
    lines_from_1 = lines_from_4 = None
    for i in sorted(inputs, key=lambda s: len(s)):
        if len(i) == 2:  # 1
            lines_from_1 = set(i)
        elif len(i) == 4:  # 4
            lines_from_4 = set(i)
        elif len(i) > 4:
            break

    lines_from_4_minus_1 = lines_from_4 - lines_from_1
    value = []
    for i in output_str.split():
        lines_from_i = set(i)
        if len(i) == 2:
            value.append("1")
        elif len(i) == 3:
            value.append("7")
        elif len(i) == 4:
            value.append("4")
        elif len(i) == 5:  # 2, 3, 5
            if not lines_from_1 - lines_from_i:
                value.append("3")
            elif not lines_from_4_minus_1 - lines_from_i:
                value.append("5")
            else:
                value.append("2")
        elif len(i) == 6:  # 6, 9 or 0
            if lines_from_4_minus_1 - lines_from_i:
                value.append("0")
            elif not lines_from_1 - lines_from_i:
                value.append("9")
            else:
                value.append("6")
        elif len(i) == 7:
            value.append("8")

    return int("".join(value))


class Day(AOCDay):
    test_solutions_p1 = [26]
    test_solutions_p2 = [61229]

    def part1(self) -> Any:
        signals = self.getInputAsArraySplit(" | ")
        count = sum(len(o) in [2, 3, 4, 7] for x in signals for o in x[1].split())

        return count

    def part2(self) -> Any:
        signals = self.getInputAsArraySplit(" | ")
        count = sum(getDisplayValue(x[0], x[1]) for x in signals)

        return count