import os
import torch

from Games.TicTacToe.TicTacToe import TicTacToe
from Games.TicTacToe.TicTacToeNN import ResNet
from Alpha_Zero_Parallel import Alpha_Zero

GAME = "TicTacToe"

args = {
    "MODEL_PATH" : os.path.join(os.getcwd(), "Games", GAME, "models_n_optimizers"),
    "SAVE_GAME_PATH" :  os.path.join(os.getcwd(), "Games", GAME, "games"),

    "EXPLORATION_CONSTANT" : 2,

    "TEMPERATURE" : 2,

    "DIRICHLET_EPSILON" : 0.25,
    "DIRICHLET_ALPHA" : 0.3,
    "ROOT_RANDOMNESS": True,

    "ADVERSARIAL" : True,

    "NO_OF_SEARCHES" : 800,
    "NO_ITERATIONS" : 1,
    "SELF_PLAY_ITERATIONS" : 3000,
    "PARALLEL_PROCESS" : 500,
    "EPOCHS" : 1,
    "BATCH_SIZE" : 128,
    "MODEL_CHECK_GAMES" : 200,
    "WIN_RATIO_FOR_SAVING": 0.5,
}


game = TicTacToe()
torch.backends.cudnn.enabled = False

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

print(device, "in use")

model = ResNet(game, 9, 128, device)

optimizer = torch.optim.Adam(model.parameters(), lr=0.001, weight_decay = 0.0001)

state = game.initialise_state()

alpha_zero = Alpha_Zero(game, args, model, optimizer)

alpha_zero.learn()