#!/usr/bin/env python

import numpy as np

def bingo(board):
    for i in range(5):
        if np.all(board[:, i] == 1) or np.all(board[i, :] == 1):
            return True
    return False

def any_bingo(boards):
    for index, board in enumerate(boards):
        if bingo(board):
            yield index
    yield -1

with open("day4.txt") as data:
    order = next(data).strip().split(',')
    next(data)
    boards = []
    current_board = []
    for line in data:
        if line == '\n':
            boards.append(current_board)
            current_board = []
        else:
            current_board.append(line.strip().split())
    boards = np.array(boards).astype(int)
    print(boards.shape)
    marked = np.zeros_like(boards)
    # part 1
    for num in map(int, order):
        indices = np.where(boards == num)
        marked[indices] = 1
        if (board_num := next(any_bingo(marked))) != -1:
            ind = np.where(marked[board_num] == 0)
            total = np.sum(boards[board_num][ind])
            print(total * num)
            break
    # part 2
    marked[:] = 0
    won = None
    for num in map(int, order):
        indices = np.where(boards == num)
        marked[indices] = 1
        nextwon = {index for index in any_bingo(marked) if index != -1}
        if len(nextwon) == 99:
            diff = (nextwon - won).pop()
            print(diff)
            ind = np.where(marked[diff] == 0)
            total = np.sum(boards[diff][ind])
            print(total*num)
            break
        else:
            won = nextwon