def main():
    import numpy as np
    from collections import defaultdict
    with open("/home/yassinealouini/Documents/code/advent_of_code/aoc/year_2021/data/4.txt") as f:
        tmp = f.read().rstrip().split("\n")

        draws = [int(e) for e in tmp[0].split(",")]

        # Make the bingo boards
        boards = defaultdict(lambda: np.zeros((5, 5)))
        board_id = 0
        data =  tmp[2:]
        for i, row in enumerate(data):
            if (i + 1) % 6 == 0:
                board_id += 1
                continue
            boards[board_id][i%6, :] = [int(e) for e in row.split(" ") if e != ""]
    
    board_won = {k: 0 for k in boards.keys()}
    for draw in draws:
        for board_id, board in boards.items():  
            board = np.where(board == draw, -1, board)
            boards[board_id] = board
            if (board.sum(axis=1) == -5).any() or (board.sum(axis=0) == -5).any():
                if sum(board_won.values())  == 0:
                    print("Solution for part I: ", board[(board > -1)].sum() * draw)
                board_won[board_id] = 1

            if sum(board_won.values()) == len(boards):
                print("Solution for part II: ", board[(board > -1)].sum() * draw)
                return

def streamlit_4(data_input):
    import numpy as np
    from collections import defaultdict
    import streamlit as st
    tmp = data_input.rstrip().split("\n")

    draws = [int(e) for e in tmp[0].split(",")]

    # Make the bingo boards
    boards = defaultdict(lambda: np.zeros((5, 5)))
    board_id = 0
    data =  tmp[2:]
    for i, row in enumerate(data):
        if (i + 1) % 6 == 0:
            board_id += 1
            continue
        boards[board_id][i%6, :] = [int(e) for e in row.split(" ") if e != ""]
    
    board_won = {k: 0 for k in boards.keys()}
    for draw in draws:
        for board_id, board in boards.items():  
            board = np.where(board == draw, -1, board)
            boards[board_id] = board
            if (board.sum(axis=1) == -5).any() or (board.sum(axis=0) == -5).any():
                if sum(board_won.values())  == 0:
                    st.write("Solution for part I: ", board[(board > -1)].sum() * draw)
                board_won[board_id] = 1

            if sum(board_won.values()) == len(boards):
                st.write("Solution for part II: ", board[(board > -1)].sum() * draw)
                return                

if __name__ == "__main__":
    main()