"""
FileName: app.py
Author: Benhao Huang
Create Date: 2023/11/19
Description: this file is used to display our project and add visualization elements to the game, using Streamlit
"""

import time
import pandas as pd
from copy import deepcopy
import numpy as np
import streamlit as st
from scipy.signal import convolve  # this is used to check if any player wins
from streamlit import session_state
from streamlit_server_state import server_state, server_state_lock
from Gomoku_MCTS import MCTSpure, alphazero, Board, PolicyValueNet_old, PolicyValueNet_new, duel_PolicyValueNet, \
    Gumbel_MCTSPlayer
from Gomoku_Bot import Gomoku_bot
from Gomoku_Bot import Board as Gomoku_bot_board

import matplotlib.pyplot as plt


from const import (
    _BLACK,  # 1
    _WHITE,  # 2
    _HUMAN,
    _AI,
    _BLANK,
    _PLAYER_NAME,
    _PLAYER_SYMBOL1,
    _PLAYER_SYMBOL2,
    _ROOM_COLOR,
    _VERTICAL,
    _NEW,
    _HORIZONTAL,
    _DIAGONAL_UP_LEFT,
    _DIAGONAL_UP_RIGHT,
    _BOARD_SIZE,
    _MODEL_PATH
)

_PLAYER_SYMBOL = [0, _PLAYER_SYMBOL1, _PLAYER_SYMBOL2]


if "FirstPlayer" not in session_state:
    session_state.FirstPlayer = _HUMAN
    session_state.Players = [ _BLACK,_WHITE]
    session_state.Symbols = _PLAYER_SYMBOL1

# Utils
class Room:
    def __init__(self, room_id) -> None:
        self.ROOM_ID = room_id
        self.BOARD = Board(width=_BOARD_SIZE, height=_BOARD_SIZE, n_in_row=5, players=session_state.Players)
        self.TURN = _BLACK
        self.CURR_PLAYER = session_state.FirstPlayer
        self.HISTORY = (0, 0)
        self.WINNER = _BLANK
        self.TIME = time.time()
        self.gomoku_bot_board = Gomoku_bot_board(_BOARD_SIZE, 1)
        self.MCTS_dict = {'Pure MCTS': MCTSpure(c_puct=5, n_playout=1000),
                          'AlphaZero': alphazero(PolicyValueNet_new(_BOARD_SIZE, _BOARD_SIZE,
                                                                    _MODEL_PATH["AlphaZero"]).policy_value_fn,
                                                 c_puct=5, n_playout=100),
                          'duel': alphazero(duel_PolicyValueNet(_BOARD_SIZE, _BOARD_SIZE,
                                                                _MODEL_PATH["duel"]).policy_value_fn,
                                            c_puct=5, n_playout=100),
                          'Gumbel AlphaZero': Gumbel_MCTSPlayer(PolicyValueNet_new(_BOARD_SIZE, _BOARD_SIZE,
                                                                                   _MODEL_PATH["Gumbel AlphaZero"]).policy_value_fn,
                                                                c_puct=5, n_playout=100, m_action=8),
                          'Gomoku Bot': Gomoku_bot(self.gomoku_bot_board, -1)}
        self.MCTS = self.MCTS_dict['AlphaZero']
        self.last_mcts = self.MCTS
        self.AID_MCTS = self.MCTS_dict['AlphaZero']
        self.COORDINATE_1D = []
        self.current_move = -1
        self.ai_simula_time_list = []
        self.human_simula_time_list = []


def change_turn(cur):
    if cur in [_HUMAN, _AI]:
        return _HUMAN if cur == _AI else _AI
    return cur % 2 + 1


# Initialize the game
if "ROOM" not in session_state:
    session_state.ROOM = Room("local")
if "OWNER" not in session_state:
    session_state.OWNER = False
if "USE_AIAID" not in session_state:
    session_state.USE_AIAID = False

# Check server health
if "ROOMS" not in server_state:
    with server_state_lock["ROOMS"]:
        server_state.ROOMS = {}


def handle_oppo_model_selection():
    new_mct = session_state.ROOM.MCTS_dict[st.session_state['selected_oppo_model']]
    session_state.ROOM.MCTS = new_mct
    session_state.ROOM.last_mcts = new_mct
    session_state.ROOM.BOARD = Board(width=_BOARD_SIZE, height=_BOARD_SIZE, n_in_row=5)
    session_state.ROOM.gomoku_bot_board = Gomoku_bot_board(_BOARD_SIZE, 1)
    session_state.ROOM.CURR_PLAYER = session_state.FirstPlayer
    session_state.ROOM.TURN = _BLACK
    session_state.ROOM.WINNER = _BLANK  # 0
    session_state.ROOM.ai_simula_time_list = []
    session_state.ROOM.human_simula_time_list = []
    session_state.ROOM.COORDINATE_1D = []
    return


def handle_aid_model_selection():
    if st.session_state['selected_aid_model'] == 'None':
        session_state.USE_AIAID = False
        return
    session_state.USE_AIAID = True
    TreeNode = session_state.ROOM.MCTS.mcts._root  # use the same tree node
    new_mct = session_state.ROOM.MCTS_dict[st.session_state['selected_aid_model']]
    new_mct.mcts._root = deepcopy(TreeNode)
    session_state.ROOM.AID_MCTS = new_mct
    return


if 'selected_oppo_model' not in st.session_state:
    st.session_state['selected_oppo_model'] = 'AlphaZero'  # 默认值

if 'selected_aid_model' not in st.session_state:
    st.session_state['selected_aid_model'] = 'AlphaZero'  # 默认值

# Layout
TITLE = st.empty()
Model_Switch = st.empty()

TITLE.header("🤖 AI 3603 Gomoku")
selected_oppo_option = Model_Switch.selectbox('Select Opponent Model',
                                              ['Pure MCTS', 'AlphaZero', 'Gomoku Bot', 'duel', 'Gumbel AlphaZero'],
                                              index=1, key='oppo_model')

if st.session_state['selected_oppo_model'] != selected_oppo_option:
    st.session_state['selected_oppo_model'] = selected_oppo_option
    handle_oppo_model_selection()

ROUND_INFO = st.empty()
st.markdown("<br>", unsafe_allow_html=True)
BOARD_PLATE = [
    [cell.empty() for cell in st.columns([1 for _ in range(_BOARD_SIZE)])] for _ in range(_BOARD_SIZE)
]
LOG = st.empty()

# Sidebar
SCORE_TAG = st.sidebar.empty()
SCORE_PLATE = st.sidebar.columns(2)
# History scores
SCORE_TAG.subheader("Scores")

PLAY_MODE_INFO = st.sidebar.container()
MULTIPLAYER_TAG = st.sidebar.empty()
with st.sidebar.container():
    ANOTHER_ROUND = st.empty()
    RESTART = st.empty()
    GIVEIN = st.empty()
    CHANGE_PLAYER = st.empty()
    AIAID = st.empty()
    EXIT = st.empty()
selected_aid_option = AIAID.selectbox('Select Assistant Model', ['None', 'Pure MCTS', 'AlphaZero'], index=0,
                                      key='aid_model')
if st.session_state['selected_aid_model'] != selected_aid_option:
    st.session_state['selected_aid_model'] = selected_aid_option
    handle_aid_model_selection()

GAME_INFO = st.sidebar.container()
message = st.empty()
PLAY_MODE_INFO.write("---\n\n**You are Black, AI agent is White.**")
GAME_INFO.markdown(
    """
    ---
    # <span style="color:black;">Freestyle Gomoku game. 🎲</span>
    - FixedModel means you are not allowed to change model during a game
    - LeaderBoard is still in development
    - no restrictions 🚫
    - no regrets 😎
    Powered by an AlphaZero approach with our own improvements! 🚀 For the specific details, please check out our <a href="insert_report_link_here" style="color:blue;">report</a>.
    ##### Adapted and improved by us! 🌟  <a href="https://github.com/Lijiaxin0111/AI_3603_BIGHOME" style="color:blue;">Our Github repo</a>
    """,
    unsafe_allow_html=True,
)


def restart() -> None:
    """
    Restart the game.
    """
    session_state.ROOM = Room(session_state.ROOM.ROOM_ID)
    st.session_state['selected_oppo_model'] = 'AlphaZero'

def givein() -> None:
    """
    Give in to AI.
    """
    session_state.ROOM = deepcopy(session_state.ROOM)
    session_state.ROOM.WINNER = _AI
    # add 1 score to AI
    session_state.ROOM.HISTORY = (
        session_state.ROOM.HISTORY[0]
        + int(session_state.ROOM.WINNER == _AI),
        session_state.ROOM.HISTORY[1]
        + int(session_state.ROOM.WINNER == _HUMAN),
    )
    session_state.ROOM.BOARD = Board(width=_BOARD_SIZE, height=_BOARD_SIZE, n_in_row=5)
    session_state.ROOM.gomoku_bot_board = Gomoku_bot_board(_BOARD_SIZE, 1)
    session_state.ROOM.MCTS_dict = {'Pure MCTS': MCTSpure(c_puct=5, n_playout=1000),
                                    'AlphaZero': alphazero(PolicyValueNet_new(_BOARD_SIZE, _BOARD_SIZE,
                                                                              _MODEL_PATH["AlphaZero"]).policy_value_fn,
                                                           c_puct=5, n_playout=100),
                                    'duel': alphazero(duel_PolicyValueNet(_BOARD_SIZE, _BOARD_SIZE,
                                                                          _MODEL_PATH["duel"]).policy_value_fn,
                                                      c_puct=5, n_playout=100),
                                    'Gumbel AlphaZero': Gumbel_MCTSPlayer(PolicyValueNet_new(_BOARD_SIZE, _BOARD_SIZE,
                                                                                             _MODEL_PATH[
                                                                                                 "Gumbel AlphaZero"]).policy_value_fn,
                                                                          c_puct=5, n_playout=100, m_action=8),
                                    'Gomoku Bot': Gomoku_bot(session_state.ROOM.gomoku_bot_board, -1)}

    session_state.ROOM.MCTS = session_state.ROOM.MCTS_dict[st.session_state['selected_oppo_model']]
    session_state.ROOM.last_mcts = session_state.ROOM.MCTS
    session_state.ROOM.TURN = _BLACK
    session_state.ROOM.CURR_PLAYER = session_state.FirstPlayer
    session_state.ROOM.WINNER = _BLANK  # 0
    session_state.ROOM.ai_simula_time_list = []
    session_state.ROOM.human_simula_time_list = []
    session_state.ROOM.COORDINATE_1D = []

def swap_players() -> None:
    session_state.update(
        FirstPlayer=change_turn(session_state.FirstPlayer),
    )
    """
    session_state.FirstPlayer = _HUMAN
    session_state.Players = [ _BLACK,_WHITE]
    session_state.Symbols = _PLAYER_SYMBOL1
    """

    session_state.ROOM.BOARD = Board(width=_BOARD_SIZE, height=_BOARD_SIZE, n_in_row=5, players=session_state.Players)
    session_state.ROOM.gomoku_bot_board = Gomoku_bot_board(_BOARD_SIZE, 1)
    session_state.ROOM.MCTS_dict = {'Pure MCTS': MCTSpure(c_puct=5, n_playout=1000),
                                    'AlphaZero': alphazero(PolicyValueNet_new(_BOARD_SIZE, _BOARD_SIZE,
                                                                              _MODEL_PATH["AlphaZero"]).policy_value_fn,
                                                           c_puct=5, n_playout=100),
                                    'duel': alphazero(duel_PolicyValueNet(_BOARD_SIZE, _BOARD_SIZE,
                                                                          _MODEL_PATH["duel"]).policy_value_fn,
                                                      c_puct=5, n_playout=100),
                                    'Gumbel AlphaZero': Gumbel_MCTSPlayer(PolicyValueNet_new(_BOARD_SIZE, _BOARD_SIZE,
                                                                                             _MODEL_PATH[
                                                                                                 "Gumbel AlphaZero"]).policy_value_fn,
                                                                          c_puct=5, n_playout=100, m_action=8),
                                    'Gomoku Bot': Gomoku_bot(session_state.ROOM.gomoku_bot_board, -1)}

    session_state.ROOM.MCTS = session_state.ROOM.MCTS_dict[st.session_state['selected_oppo_model']]
    session_state.ROOM.last_mcts = session_state.ROOM.MCTS
    session_state.ROOM.TURN = _BLACK
    session_state.ROOM.CURR_PLAYER = session_state.FirstPlayer
    session_state.ROOM.WINNER = _BLANK  # 0
    session_state.ROOM.ai_simula_time_list = []
    session_state.ROOM.human_simula_time_list = []
    session_state.ROOM.COORDINATE_1D = []

RESTART.button(
    "Reset",
    on_click=restart,
    help="Clear the board as well as the scores",
)

GIVEIN.button(
    "Give in",
    on_click = givein,
    help="Give in to AI",
)

CHANGE_PLAYER.button(
    "Swap players",
    on_click=swap_players,
    help="Swap players",
)


# Draw the board
def gomoku():
    """
    Draw the board.
    Handle the main logic.
    """

    # Restart the game

    # Continue new round
    def another_round() -> None:
        """
        Continue new round.
        """
        session_state.ROOM = deepcopy(session_state.ROOM)
        session_state.ROOM.BOARD = Board(width=_BOARD_SIZE, height=_BOARD_SIZE, n_in_row=5)
        session_state.ROOM.gomoku_bot_board = Gomoku_bot_board(_BOARD_SIZE, 1)
        session_state.ROOM.MCTS_dict = {'Pure MCTS': MCTSpure(c_puct=5, n_playout=1000),
                                        'AlphaZero': alphazero(PolicyValueNet_new(_BOARD_SIZE, _BOARD_SIZE,
                                                                                  _MODEL_PATH["AlphaZero"]).policy_value_fn,
                                                               c_puct=5, n_playout=100),
                                        'duel': alphazero(duel_PolicyValueNet(_BOARD_SIZE, _BOARD_SIZE,
                                                                              _MODEL_PATH["duel"]).policy_value_fn,
                                                          c_puct=5, n_playout=100),
                                        'Gumbel AlphaZero': Gumbel_MCTSPlayer(PolicyValueNet_new(_BOARD_SIZE, _BOARD_SIZE,
                                                                                   _MODEL_PATH["Gumbel AlphaZero"]).policy_value_fn,
                                                                c_puct=5, n_playout=100, m_action=8),
                                        'Gomoku Bot': Gomoku_bot(session_state.ROOM.gomoku_bot_board, -1)}

        session_state.ROOM.MCTS = session_state.ROOM.MCTS_dict[st.session_state['selected_oppo_model']]
        session_state.ROOM.last_mcts = session_state.ROOM.MCTS
        session_state.ROOM.TURN = _BLACK
        session_state.ROOM.CURR_PLAYER = session_state.FirstPlayer
        session_state.ROOM.WINNER = _BLANK  # 0
        session_state.ROOM.ai_simula_time_list = []
        session_state.ROOM.human_simula_time_list = []
        session_state.ROOM.COORDINATE_1D = []

    # Room status sync
    def sync_room() -> bool:
        room_id = session_state.ROOM.ROOM_ID
        if room_id not in server_state.ROOMS.keys():
            session_state.ROOM = Room("local")
            return False
        elif server_state.ROOMS[room_id].TIME == session_state.ROOM.TIME:
            return False
        elif server_state.ROOMS[room_id].TIME < session_state.ROOM.TIME:
            # Only acquire the lock when writing to the server state
            with server_state_lock["ROOMS"]:
                server_rooms = server_state.ROOMS
                server_rooms[room_id] = session_state.ROOM
                server_state.ROOMS = server_rooms
            return True
        else:
            session_state.ROOM = server_state.ROOMS[room_id]
            return True

    # Triggers the board response on click
    def handle_click(x, y):
        """
        Controls whether to pass on / continue current board / may start new round
        """
        if session_state.ROOM.BOARD.board_map[x][y] != _BLANK:
            pass
        elif (
                session_state.ROOM.ROOM_ID in server_state.ROOMS.keys()
                and _ROOM_COLOR[session_state.OWNER]
                != server_state.ROOMS[session_state.ROOM.ROOM_ID].TURN
        ):
            sync_room()

        # normal play situation
        elif session_state.ROOM.WINNER == _BLANK:
            move = session_state.ROOM.BOARD.location_to_move((x, y))
            session_state.ROOM.current_move = move
            session_state.ROOM.BOARD.do_move(move)
            # Gomoku Bot BOARD
            session_state.ROOM.MCTS_dict["Gomoku Bot"].board.put(_BOARD_SIZE - move // _BOARD_SIZE - 1,
                                                                 move % _BOARD_SIZE)  # # this move starts from left up corner (0,0), however, the move in the game starts from left bottom corner (0,0)
            session_state.ROOM.BOARD.board_map[x][y] = session_state.ROOM.TURN
            session_state.ROOM.COORDINATE_1D.append(x * _BOARD_SIZE + y)

            session_state.ROOM.TURN = change_turn(session_state.ROOM.TURN)
            session_state.ROOM.CURR_PLAYER = change_turn(session_state.ROOM.CURR_PLAYER)
            win, winner = session_state.ROOM.BOARD.game_end()
            if win:
                session_state.ROOM.WINNER = session_state.ROOM.CURR_PLAYER
            session_state.ROOM.HISTORY = (
                session_state.ROOM.HISTORY[0]
                + int(session_state.ROOM.WINNER == _AI),
                session_state.ROOM.HISTORY[1]
                + int(session_state.ROOM.WINNER == _HUMAN),
            )
            session_state.ROOM.TIME = time.time()

    def forbid_click(x, y):
        # st.warning('This posistion has been occupied!!!!', icon="⚠️")
        st.error("({}, {}) has been occupied!!)".format(x, y), icon="🚨")

    # Draw board
    def draw_board(response: bool):
        """construct each buttons for all cells of the board"""
        if session_state.USE_AIAID and session_state.ROOM.WINNER == _BLANK and session_state.ROOM.CURR_PLAYER == _HUMAN:
            if session_state.USE_AIAID:
                copy_mcts = deepcopy(session_state.ROOM.AID_MCTS.mcts)
                _, acts_aid, probs_aid, simul_mean_time_aid = copy_mcts.get_move_probs(session_state.ROOM.BOARD)
                sorted_acts_probs = sorted(zip(acts_aid, probs_aid), key=lambda x: x[1], reverse=True)
                top_five_acts = [act for act, prob in sorted_acts_probs[:5]]
                top_five_probs = [prob for act, prob in sorted_acts_probs[:5]]
        if response and session_state.ROOM.CURR_PLAYER == _HUMAN:  # human turn
            start_time = time.time()
            print("Your turn")
            # construction of clickable buttons
            cur_move = (session_state.ROOM.current_move // _BOARD_SIZE, session_state.ROOM.current_move % _BOARD_SIZE)
            for i, row in enumerate(session_state.ROOM.BOARD.board_map):
                for j, cell in enumerate(row):
                    if (
                            i * _BOARD_SIZE + j
                            in (session_state.ROOM.COORDINATE_1D)
                    ):
                        if i == cur_move[0] and j == cur_move[1]:
                            BOARD_PLATE[i][j].button(
                                session_state.Symbols[_NEW],
                                key=f"{i}:{j}",
                                args=(i, j),
                                on_click=forbid_click,
                            )
                        else:
                            # disable click for GPT choices
                            BOARD_PLATE[i][j].button(
                                session_state.Symbols[cell],
                                key=f"{i}:{j}",
                                args=(i, j),
                                on_click=forbid_click
                            )
                    else:
                        if session_state.USE_AIAID and i * _BOARD_SIZE + j in top_five_acts:
                            # enable click for other cells available for human choices
                            prob = top_five_probs[top_five_acts.index(i * _BOARD_SIZE + j)]
                            BOARD_PLATE[i][j].button(
                                session_state.Symbols[cell] + f"{round(prob, 2)}",
                                key=f"{i}:{j}",
                                on_click=handle_click,
                                args=(i, j),
                            )
                        else:
                            # enable click for other cells available for human choices
                            BOARD_PLATE[i][j].button(
                                session_state.Symbols[cell],
                                key=f"{i}:{j}",
                                on_click=handle_click,
                                args=(i, j),
                            )
            end_time = time.time()
            print("Time used for human move: ", end_time - start_time)

        elif response and session_state.ROOM.CURR_PLAYER == _AI:  # AI turn
            message.empty()
            with st.spinner('🔮✨ Waiting for AI response... ⏳🚀'):
                time.sleep(0.05)
                print("AI's turn")
                
                if st.session_state['selected_oppo_model'] != 'Gomoku Bot':
                    move, simul_time = session_state.ROOM.MCTS.get_action(session_state.ROOM.BOARD, return_time=True)
                else:
                    move, simul_time = session_state.ROOM.MCTS.get_action(return_time=True)
                session_state.ROOM.ai_simula_time_list.append(simul_time)
                print("AI takes move: ", move)
                session_state.ROOM.current_move = move
                gpt_response = move
                gpt_i, gpt_j = gpt_response // _BOARD_SIZE, gpt_response % _BOARD_SIZE
                print("AI's move is located at ({}, {}) :".format(gpt_i, gpt_j))
                move = session_state.ROOM.BOARD.location_to_move((gpt_i, gpt_j))
                print("Location to move: ", move)
                # MCTS BOARD
                session_state.ROOM.BOARD.do_move(move)
                # Gomoku Bot BOARD
                session_state.ROOM.MCTS_dict["Gomoku Bot"].board.put(_BOARD_SIZE - 1 - move // _BOARD_SIZE,
                                                                     move % _BOARD_SIZE)
                session_state.ROOM.COORDINATE_1D.append(gpt_i * _BOARD_SIZE + gpt_j)

                if not session_state.ROOM.BOARD.game_end()[0]:
                    if session_state.USE_AIAID:
                        copy_mcts = deepcopy(session_state.ROOM.AID_MCTS.mcts)
                        _, acts_aid, probs_aid, simul_mean_time_aid = copy_mcts.get_move_probs(session_state.ROOM.BOARD)
                        sorted_acts_probs = sorted(zip(acts_aid, probs_aid), key=lambda x: x[1], reverse=True)
                        top_five_acts = [act for act, prob in sorted_acts_probs[:5]]
                        top_five_probs = [prob for act, prob in sorted_acts_probs[:5]]
                else:
                    top_five_acts = []
                    top_five_probs = []

                # construction of clickable buttons
                for i, row in enumerate(session_state.ROOM.BOARD.board_map):
                    # print("row:", row)
                    for j, cell in enumerate(row):
                        if (
                                i * _BOARD_SIZE + j
                                in (session_state.ROOM.COORDINATE_1D)
                        ):
                            if i == gpt_i and j == gpt_j:
                                BOARD_PLATE[i][j].button(
                                    session_state.Symbols[_NEW],
                                    key=f"{i}:{j}",
                                    args=(i, j),
                                    on_click=handle_click,
                                )
                            else:
                                # disable click for GPT choices
                                BOARD_PLATE[i][j].button(
                                    session_state.Symbols[cell],
                                    key=f"{i}:{j}",
                                    args=(i, j),
                                    on_click=forbid_click
                                )
                        else:
                            if session_state.USE_AIAID and i * _BOARD_SIZE + j in top_five_acts and not \
                                    session_state.ROOM.BOARD.game_end()[0]:
                                # enable click for other cells available for human choices
                                prob = top_five_probs[top_five_acts.index(i * _BOARD_SIZE + j)]
                                BOARD_PLATE[i][j].button(
                                    session_state.Symbols[cell] + f"{round(prob, 2)}",
                                    key=f"{i}:{j}",
                                    on_click=handle_click,
                                    args=(i, j),
                                )
                            else:
                                # enable click for other cells available for human choices
                                BOARD_PLATE[i][j].button(
                                    session_state.Symbols[cell],
                                    key=f"{i}:{j}",
                                    on_click=handle_click,
                                    args=(i, j),
                                )

            message.markdown(
                'AI agent has calculated its strategy, which takes <span style="color: blue; font-size: 20px;">{:.3e}</span>s per simulation.'.format(
                    simul_time),
                unsafe_allow_html=True
            )
            LOG.subheader("Logs")
            # change turn
            session_state.ROOM.TURN = change_turn(session_state.ROOM.TURN)
            session_state.ROOM.CURR_PLAYER = change_turn(session_state.ROOM.CURR_PLAYER)

            win, winner = session_state.ROOM.BOARD.game_end()
            if win:
                session_state.ROOM.WINNER = session_state.ROOM.CURR_PLAYER

            session_state.ROOM.HISTORY = (
                session_state.ROOM.HISTORY[0]
                + int(session_state.ROOM.WINNER == _AI),
                session_state.ROOM.HISTORY[1]
                + int(session_state.ROOM.WINNER == _HUMAN),
            )
            session_state.ROOM.TIME = time.time()

        if not response or session_state.ROOM.WINNER != _BLANK:
            if session_state.ROOM.WINNER != _BLANK:
                print("Game over")
            for i, row in enumerate(session_state.ROOM.BOARD.board_map):
                for j, cell in enumerate(row):
                    BOARD_PLATE[i][j].write(
                        session_state.Symbols[cell],
                        # key=f"{i}:{j}",
                    )

    # Game process control
    def game_control():
        if session_state.ROOM.WINNER != _BLANK:
            draw_board(False)
        else:
            draw_board(True)
        if session_state.ROOM.WINNER != _BLANK or 0 not in session_state.ROOM.BOARD.board_map:
            GIVEIN.empty()
            ANOTHER_ROUND.button(
                "Play Next round!",
                on_click=another_round,
                help="Clear board and swap first player",
            )

    # Infos
    def update_info() -> None:
        # Additional information
        SCORE_PLATE[0].metric("Gomoku-Agent", session_state.ROOM.HISTORY[0])
        SCORE_PLATE[1].metric("You", session_state.ROOM.HISTORY[1])
        if session_state.ROOM.WINNER != _BLANK:
            st.balloons()
            ROUND_INFO.write(
                f"#### **{_PLAYER_NAME[session_state.ROOM.WINNER]} WIN!**\n**Click buttons on the left for more plays.**"
            )

        st.markdown("<br>", unsafe_allow_html=True)
        st.markdown("<br>", unsafe_allow_html=True)
        chart_data = pd.DataFrame(session_state.ROOM.ai_simula_time_list, columns=["Simulation Time"])
        st.line_chart(chart_data)

    game_control()
    update_info()


if __name__ == "__main__":
    gomoku()