from copy import copy
from dataclasses import dataclass
from typing import List, Tuple, Callable

from aoc import BaseAssignment


def bold(item: str, condition: Callable[[], bool]):
    return f'\033[36m{item}\033[0m' if condition() else item

@dataclass
class Number:
    nr: int
    marked: bool = False

    def __int__(self):
        return self.nr

    def __bool__(self):
        return self.marked

    def __radd__(self, other):
        return other + int(self)

    def mark(self):
        self.marked = True

@dataclass
class BingoCard:
    card: List[List[Number]]

    @property
    def bingo(self) -> bool:
        return any([
            *[all(bool(item) for item in row) for row in self.card],
            *[all(bool(row[col_index]) for row in self.card) for col_index in range(5)]
        ])

    def mark(self, nr: int):
        for row in self.card:
            for item in row:
                if int(item) == nr:
                    item.mark()

    @property
    def unmarked(self):
        return [
            item
            for row in self.card
            for item in row
            if not bool(item)
        ]

    def __str__(self):
            return '\n'.join([
                ''.join([
                    bold(str(int(item)).rjust(2, ' ').ljust(3, ' '), lambda: item.marked)
                    for item in row
                ])
                for row in self.card
            ])


class Assignment(BaseAssignment):
    def read_input(self, example = False) -> Tuple[List[int], List[BingoCard]]:
        cards = []
        card = None

        input = super().read_input(example)
        nrs = [ int(n) for n in next(input).split(',') ]

        try:
            row = next(input)

            while True:
                if row == '':
                    if card is not None:
                        cards.append(BingoCard(card=card))
                    card = []
                else:
                    card.append([ Number(nr=int(item)) for item in row.split(' ') if item != ''])

                row = next(input)

        except StopIteration:
            if card is not None:
                cards.append(BingoCard(card=card))

        return nrs, cards


class AssignmentOne(Assignment):
    example_result = 4512

    def run(self, input: Tuple[List[int], List[BingoCard]]) -> int:
        nrs, cards = input

        for nr in nrs:
            for card in cards:
                card.mark(nr)

                if card.bingo:
                    return nr * sum(card.unmarked)


class AssignmentTwo(Assignment):
    example_result = 1924

    def run(self, input: Tuple[List[int], List[BingoCard]]) -> int:
        nrs, cards = input

        in_game = copy(cards)
        for nr in nrs:
            for card in copy(in_game):
                card.mark(nr)

                if card.bingo:
                    in_game.remove(card)

                if len(in_game) == 0:
                    return nr * sum(card.unmarked)