import streamlit as st
import numpy as np
from objects_9x9 import Player, HumanPlayer

# Create RL bot player and human player
p1 = Player("9x9_p1")
human_player = HumanPlayer("p2")

def handle_click(i, j):
    if (i, j) not in check_available_moves(extra=True):
        st.session_state.warning = True
    elif not st.session_state.winner:
        st.session_state.warning = False
        st.session_state.board[i, j] = st.session_state.player
        winner = check_win(st.session_state.board)
        if winner:
            st.session_state.winner = winner
            if st.session_state.opponent == 'Computer':
                # Give reward to the RL bot and update its policy
                if winner == 'X':
                    p1.feedReward(1)
                elif winner == 'O':
                    p1.feedReward(0)
                else:
                    p1.feedReward(0.1)
                
                # Save the RL bot's policy
                p1.savePolicy()
        
        # Toggle the player's turn
        st.session_state.player = "O" if st.session_state.player == "X" else "X"

def init(post_init=False):
    if not post_init:
        st.session_state.win = {'X': 0, 'O': 0}
        st.session_state.opponent = 'Computer'  # Initialize the opponent attribute
    st.session_state.board = np.full((9, 9), '.', dtype=str)
    st.session_state.player = 'X'
    st.session_state.warning = False
    st.session_state.winner = None
    st.session_state.over = False

def check_available_moves(extra=False) -> list:
    raw_moves = [row for col in st.session_state.board.tolist() for row in col]
    num_moves = [i for i, spot in enumerate(raw_moves) if spot == '.']
    if extra:
        return [(i // 9, i % 9) for i in num_moves]
    return num_moves

def check_win(board):
    # Check for a winner by checking rows, columns, and diagonals for three linked points
    for i in range(9):
        for j in range(9):
            if board[i, j] != '.':
                symbol = board[i, j]
                # Check horizontally
                if j + 2 < 9 and board[i, j + 1] == board[i, j + 2] == symbol:
                    return symbol
                # Check vertically
                if i + 2 < 9 and board[i + 1, j] == board[i + 2, j] == symbol:
                    return symbol
                # Check diagonally (top-left to bottom-right)
                if i + 2 < 9 and j + 2 < 9 and board[i + 1, j + 1] == board[i + 2, j + 2] == symbol:
                    return symbol
                # Check diagonally (top-right to bottom-left)
                if i + 2 < 9 and j - 2 >= 0 and board[i + 1, j - 1] == board[i + 2, j - 2] == symbol:
                    return symbol
    return None

def computer_player():
    moves = check_available_moves(extra=True)
    if moves:
        # Use p1 to choose the action
        positions = check_available_moves(extra=True)
        p1_action = p1.chooseAction(positions, st.session_state.board, -1)
        
        # Check if the chosen action is valid and make the move
        if p1_action in moves:
            i, j = p1_action
            handle_click(i, j)

def main():
    st.write(
        """
        # ❎🅾️ Tic Tac Toe (9x9)
        """
    )

    if "board" not in st.session_state:
        init()

    reset, score, player = st.columns([0.5, 0.6, 1])
    reset.button('New game', on_click=init, args=(True,))

    for i, row in enumerate(st.session_state.board):
        cols = st.columns([5] + [1] * 9 + [5])
        for j, field in enumerate(row):
            if st.session_state.player == 'X' and not st.session_state.winner:
                # Only allow human player to make a move if it's their turn and the game is not over
                cols[j + 1].button(
                    field,
                    key=f"{i}-{j}",
                    on_click=lambda i=i, j=j: handle_click(i, j)  # Use lambda to pass arguments
                )
            else:
                # If it's not the human player's turn or the game is over, let the bot play
                computer_player()

    if st.session_state.winner:
        st.success(f"Congrats! {st.session_state.winner} won the game! 🎈")
    elif not check_available_moves():
        st.info(f"It's a tie 📍")

    score.button(f'❌{st.session_state.win["X"]} 🆚 {st.session_state.win["O"]}⭕')
    player.button(
        f'{"❌" if st.session_state.player == "X" else "⭕"}\'s turn'
        if not st.session_state.winner
        else f'🏁 Game finished'
    )

if __name__ == '__main__':
    main()