from __future__ import annotations

from pathlib import Path

import gradio as gr
import numpy as np
from game import is_terminal, legal_actions, play, winner
from mcts import NeuralMCTS
from model import MicroZeroNet
from safetensors.torch import load_file

ARTIFACT = Path(__file__).resolve().parent / "artifacts" / "microzero"
MODEL = MicroZeroNet()
MODEL.load_state_dict(load_file(ARTIFACT / "model.safetensors"))
MODEL.eval()
SYMBOLS = {0: ".", 1: "X", -1: "O"}


def neural_action(
    model: MicroZeroNet,
    board: np.ndarray,
    player: int,
    simulations: int,
    seed: int,
) -> int:
    search = NeuralMCTS(model, simulations=simulations)
    policy = search.policy(
        board,
        player,
        temperature=0.0,
        add_noise=False,
        rng=np.random.default_rng(seed),
    )
    return int(np.argmax(policy))


def render(board: list[int], message: str) -> str:
    cells = [SYMBOLS[int(value)] for value in board]
    rows = [" | ".join(cells[index : index + 3]) for index in range(0, 9, 3)]
    return "```\n" + "\n---------\n".join(rows) + f"\n```\n{message}"


def new_game() -> tuple[list[int], str]:
    board = [0] * 9
    return board, render(board, "You are X. Choose a cell from 0 to 8.")


def move(board: list[int], action: int) -> tuple[list[int], str]:
    array = np.asarray(board, dtype=np.int8)
    action = int(action)
    if action not in legal_actions(array):
        return board, render(board, "That cell is occupied.")
    array = play(array, action, 1)
    if is_terminal(array):
        message = "You win." if winner(array) == 1 else "Draw."
        return array.tolist(), render(array.tolist(), message)
    response = neural_action(MODEL, array, -1, simulations=96, seed=7000 + action)
    array = play(array, response, -1)
    result = winner(array)
    if result == -1:
        message = f"MicroZero played {response} and won."
    elif is_terminal(array):
        message = f"MicroZero played {response}. Draw."
    else:
        message = f"MicroZero played {response}. Your turn."
    return array.tolist(), render(array.tolist(), message)


with gr.Blocks(title="MicroZero") as demo:
    gr.Markdown("# MicroZero\nPlay against the neural MCTS agent.")
    state = gr.State([0] * 9)
    board = gr.Markdown(render([0] * 9, "You are X. Choose a cell from 0 to 8."))
    action = gr.Slider(0, 8, value=4, step=1, label="Your move")
    with gr.Row():
        play_button = gr.Button("Play", variant="primary")
        reset_button = gr.Button("New game")
    play_button.click(move, inputs=[state, action], outputs=[state, board])
    reset_button.click(new_game, outputs=[state, board])


if __name__ == "__main__":
    demo.launch()