# imports
import numpy as np
import torch
import torch.nn as nn
import chess
from chess import Move

# set up AI (Sequential NN)


class Tanh200(nn.Module):
    def forward(self, x):
        return torch.tanh(x / 200)


model = nn.Sequential(
    nn.Linear(833, 1024),
    nn.BatchNorm1d(1024),
    nn.ReLU(),
    nn.Linear(1024, 512),
    nn.Dropout(p=0.25),
    nn.BatchNorm1d(512),
    nn.LeakyReLU(0.075),
    nn.Linear(512, 256),
    nn.BatchNorm1d(256),
    nn.ReLU(),
    nn.Linear(256, 128),
    nn.Dropout(p=0.25),
    nn.BatchNorm1d(128),
    nn.LeakyReLU(0.075),
    nn.Linear(128, 64),
    nn.ReLU(),
    nn.Linear(64, 1),
    nn.Dropout(p=0.25),
    Tanh200(),
)

model.eval()
weights_path = "./zlpro.pt"
state_dict = torch.load(weights_path)
model.load_state_dict(state_dict)


# set up important functions
# board_data - this function will turn a python-chess board into a matrix for the AI to evaluate
def board_data(board):
    board_array = np.zeros((8, 8, 13), dtype=np.int8)
    for i in range(64):
        piece = board.piece_at(i)
        if piece is not None:
            color = int(piece.color)
            piece_type = piece.piece_type - 1
            board_array[i // 8][i % 8][piece_type + 6 * color] = 1
        else:
            board_array[i // 8][i % 8][-1] = 1
    board_array = board_array.flatten()
    return board_array

def negamax_ab(
    board,
    alpha,
    beta,
    move_count,
    turn_to_move,
    depth=2,
):
    if depth == 0:
        move_turn = turn_to_move
        matrix_game = np.array([board_data(board)])  # game after one hot encoding
        matrix_game = np.concatenate(
            (matrix_game, np.array([[move_turn]])), axis=1
        )  # have to append the move turn - the AI needs to know this
        matrix_game = torch.tensor(matrix_game, dtype=torch.float32)
        with torch.no_grad():
            best_val = model(matrix_game)
        best_val = float(best_val)
        return best_val * turn_to_move
    child_nodes = list(board.legal_moves)
    # child_nodes = ordermoves(child_nodes) # make an ordermove function
    best_score = -np.inf
    indexer = 0
    for child in child_nodes:
        test_board = board.copy()
        test_board.push(child)
        child = board_data(test_board)
        score = negamax_ab(
            board,
            -beta,
            -alpha,
            move_count,
            turn_to_move,
            depth - 1,
        )
        score = -float(score)
        best_score = max(best_score, score)
        alpha = max(alpha, best_score)
        test_board.pop()
        if alpha >= beta:
            break
        indexer += 1
    return best_score


# set up chess game

NUMBER_OF_GAMES = 10


def play_game(NUMBER_OF_GAMES):
    for _ in range(NUMBER_OF_GAMES):
        board = chess.Board()
        legal_moves = board.legal_moves
        move_count = 0
        for __ in range(3):  # random opening, 3 moves
            legal_moves = list(board.legal_moves)
            chosen_move = np.random.randint(0, len(legal_moves))
            board.push(legal_moves[chosen_move])
            with open("ai_games.txt", "a+") as f:
                f.write(str(legal_moves[chosen_move]))
                f.write(" ")
        while not board.is_game_over():
            if move_count % 2 == 0:
                turn_to_move = 1  # white
            else:
                turn_to_move = -1  # black
            legal_moves = list(board.legal_moves)
            m_dict = {}
            for move in legal_moves:
                test_board = board.copy()
                test_board.push(move)
                move_score = negamax_ab(
                    test_board, -np.inf, np.inf, move_count, turn_to_move, 3
                )
                m_dict[str(move)] = move_score
                m_dict = {
                    k: v
                    for k, v in sorted(
                        m_dict.items(), key=lambda item: item[1], reverse=True
                    )
                }
                m = iter(m_dict)
                best_move = next(m)
            print("Scores of all moves:", m_dict)
            print("Best move:", best_move)
            with open("ai_games.txt", "a+") as f:
                f.write(best_move)
                f.write(" ")
            best_move = Move.from_uci(best_move)
            board.push(best_move)
            move_count += 1
        with open("ai_games.txt", "a+") as f:
            f.write(board.result())
            f.write("\n")


play_game(NUMBER_OF_GAMES)