from typing import Any, Iterator

import numpy as np

from ..base import BaseSolver


class Solver(BaseSolver):
    def solve(self, input: str) -> Iterator[Any]:
        lines = input.splitlines()

        numbers = [int(c) for c in lines[0].split(",")]

        boards = np.asarray(
            [
                [[int(c) for c in line.split()] for line in lines[start : start + 5]]
                for start in range(2, len(lines), 6)
            ]
        )

        # (round, score) for each board (-1 when not found)
        winning_rounds: list[tuple[int, int]] = [(-1, -1) for _ in range(len(boards))]
        marked = np.zeros_like(boards, dtype=bool)

        for round, number in enumerate(numbers):
            # mark boards
            marked[boards == number] = True

            # check each board for winning
            for index in range(len(boards)):
                if winning_rounds[index][0] > 0:
                    continue

                if np.any(
                    np.all(marked[index], axis=0) | np.all(marked[index], axis=1)
                ):
                    winning_rounds[index] = (
                        round,
                        number * int(np.sum(boards[index][~marked[index]])),
                    )

            # all boards are winning - break
            if np.all(marked.all(axis=1) | marked.all(axis=2)):
                break

        # part 1
        (_, score) = min(winning_rounds, key=lambda w: w[0])
        yield score

        # part 2
        (_, score) = max(winning_rounds, key=lambda w: w[0])
        yield score