import copy

import os

import pygame

import numpy as np

from numpy import sqrt

from pygame.locals import *

import time

import global_vars as gv

import tensorflow as tf

import architectures.tree_tf_op as arch

import py_util.py_util as pu

########################################################## configuration:

#save_nm = 'models/go_cpu_tree_0.200000EPS_7GMSZ_1000N_SIM_0.001000L2_LAMBDA_0.900000MOMENTUM_0.025000VAL_LAMBDA_1.000000CPUCT_20N_TURNS_128N_FILTERS_EPS0.110000_EPS0.020000.npy'

save_nm = 'models/go_cpu_tree_0.200000EPS_7GMSZ_1000N_SIM_0.001000L2_LAMBDA_0.900000MOMENTUM_0.025000VAL_LAMBDA_1.000000CPUCT_20N_TURNS_128N_FILTERS_EPS0.110000_EPS0.020000_EPS0.010000.npy'

# ^ set save_nm = None if you want to start training a new model

#run_net = True # run only the network (no tree search)

run_net = False # make moves from the tree search

show_txt = False # don't show statistics of each move (Q and P values, visit counts)

#show_txt = True

# load the following variables from the model .npy file:

save_vars = ['LSQ_LAMBDA', 'LSQ_REG_LAMBDA', 'POL_CROSS_ENTROP_LAMBDA', 'VAL_LAMBDA', 'VALR_LAMBDA', 'L2_LAMBDA',

'FILTER_SZS', 'STRIDES', 'N_FILTERS', 'N_FC1', 'EPS', 'MOMENTUM', 'SAVE_FREQ', 'N_SIM', 'N_TURNS', 'CPUCT', 'DIR_A']

save_d = np.load(save_nm).item()

for key in save_vars:

if key == 'save_nm':

continue

exec('%s = save_d["%s"]' % (key,key))

########## over-write number of simulations previously used:

N_SIM = 1000

CPUCT = 1

if DIR_A == 0:

dir_pre = 0

else:

dir_pre = gamma(DIR_A * gv.map_szt) / (gamma(DIR_A)**gv.map_szt)

###############################################################################

############## load model, init variables

arch.init_model(N_FILTERS, FILTER_SZS, STRIDES, N_FC1, EPS, MOMENTUM,

LSQ_LAMBDA, LSQ_REG_LAMBDA, POL_CROSS_ENTROP_LAMBDA, VAL_LAMBDA, VALR_LAMBDA, L2_LAMBDA)

arch.saver.restore(arch.sess, save_nm)

arch.sess.run(arch.init_state)

pu.init_tree()

Q_map = np.zeros((gv.n_rows, gv.n_cols), dtype='single')

Q_map_next = np.zeros_like(Q_map) # Q values for the move after the current (assuming you make the move the network predicts you will)

P_map = np.zeros_like(Q_map)

P_map_next = np.zeros_like(Q_map)

visit_count_map = np.zeros((gv.n_rows, gv.n_cols), dtype='int32')

visit_count_map_next = np.zeros_like(visit_count_map)

def ret_d(player): # return dictionary for input into tensor flow

return {arch.moving_player: player, arch.dir_a: DIR_A, arch.dir_pre: dir_pre}

def ret_stats(player): # return Q map, P map, and visit count maps

return Q_map, P_map, visit_count_map

# move neural network

def nn_mv():

return arch.sess.run(arch.gm_vars['board'])[0]

##################### display

psz = 50 # size to display pieces

pszh = psz/2.

n_txt_rows = 4

window_sz = (psz*gv.n_rows, psz*gv.n_cols)

BLACK = (0,)*3

LINE_WIDTH = 2

turn = 0

windowSurface = pygame.display.set_mode(window_sz, 0, 32)

pygame.display.set_caption('Go GUI')

pygame.init()

basicFont = pygame.font.SysFont(None, 22)

whitep = pygame.image.load('notebooks/go_white.png')

blackp = pygame.image.load('notebooks/go_black.png')

blank = pygame.image.load('notebooks/go_blank.png')

whitep = pygame.transform.scale(whitep, (psz, psz))

blackp = pygame.transform.scale(blackp, (psz, psz))

blank = pygame.transform.scale(blank, window_sz)

centers = np.arange(gv.n_rows)*psz + pszh

to_coords = -np.ones(gv.BATCH_SZ, dtype='int32')

board = np.zeros((gv.n_rows, gv.n_cols), dtype='int8')

# draw board and optionally text

def draw(update=False):

pygame.display.update()

draw(update=True)

board = nn_mv()

draw(update=True)

while True:

event = pygame.event.wait()

if event.type == MOUSEBUTTONUP:

mouse_pos = np.asarray(event.pos)

x = np.argmin((mouse_pos[0] - centers)**2)

y = np.argmin((mouse_pos[1] - centers)**2)

to_coords[0] = x*gv.n_cols + y

board_prev = arch.sess.run(arch.gm_vars['board'])[0]

imgs = arch.sess.run(arch.imgs, feed_dict={arch.moving_player: 1})

arch.sess.run(arch.nn_max_move_unit, feed_dict={arch.moving_player: 1, arch.nn_max_to_coords: to_coords})

board = arch.sess.run(arch.gm_vars['board'])[0]

if board_prev.sum() == board.sum(): # invalid move

print 'invalid mv'

continue

win_tmp, score_tmp = arch.sess.run([arch.winner, arch.score], feed_dict={arch.moving_player: 1})

print 'turn %i, winner %i, score %i' % (turn, win_tmp[0], score_tmp[0])

draw(update=True)

board = nn_mv()

draw(update=True)

turn += 1

win_tmp, score_tmp = arch.sess.run([arch.winner, arch.score], feed_dict={arch.moving_player: 1})

print 'turn %i, winner %i, score %i' % (turn, win_tmp[0], score_tmp[0])

if event.type == QUIT:

pygame.display.quit()

break