import os
import time
import numpy as np
import sys

from utils import *
from NeuralNet import NeuralNet

from .ChessNNet import ChessNNet as onnet

sys.path.append('..')

args = TrainingConfig({
    'lr': 0.005,
    'dropout': 0.2,
    'epochs': 14,
    'batch_size': 16,
    'cuda': True,
    'num_channels': 256,
})


class NNetWrapper(NeuralNet):
    def __init__(self, game):
        self.nnet = onnet(game, args)
        self.board_x, self.board_y = game.getBoardSize()
        self.action_size = game.getActionSize()

    def train(self, examples):
        """
        examples: list of examples, each example is of form (board, pi, v)
        """
        input_boards, target_pis, target_vs = list(zip(*examples))
        input_boards = np.asarray(input_boards)
        target_pis = np.asarray(target_pis)
        target_vs = np.asarray(target_vs)
        self.nnet.model.fit(x=input_boards, y=[target_pis, target_vs], batch_size=args.batch_size, epochs=args.epochs, verbose=2)

    def predict(self, board):
        """
        board: np array with board
        """
        # timing
        #start = time.time()

        # preparing input
        board = board[np.newaxis, :, :]

        # run
        pi, v = self.nnet.model.predict(board)

        #print('PREDICTION TIME TAKEN : {0:03f}'.format(time.time()-start))
        return pi[0], v[0]

    def save_checkpoint(self, folder='checkpoint', filename='checkpoint.h5'):
        filepath = os.path.join(folder, filename)
        if not os.path.exists(folder):
            print(f"Checkpoint Directory does not exist! Making directory {folder}")
            os.mkdir(folder)
        else:
            print(f"Checkpoint Directory exists (at '{os.path.abspath(filepath)}').")
        self.nnet.model.save_weights(filepath)

    def load_checkpoint(self, folder='checkpoint', filename='checkpoint.h5'):
        filepath = os.path.join(folder, filename)
        if not os.path.exists(filepath):
            print(f"No model in path '{os.path.abspath(filepath)}'!")
            raise Exception("No model in path!")
        self.nnet.model.load_weights(filepath)

import os
import time
import numpy as np
import sys

from utils import *
from NeuralNet import NeuralNet

from .ChessNNet import ChessNNet as onnet

sys.path.append('..')

args = TrainingConfig({
    'lr': 0.005,
    'dropout': 0.2,
    'epochs': 14,
    'batch_size': 16,
    'cuda': True,
    'num_channels': 256,
})


class NNetWrapper(NeuralNet):
    def __init__(self, game):
        self.nnet = onnet(game, args)
        self.board_x, self.board_y = game.getBoardSize()
        self.action_size = game.getActionSize()

    def train(self, examples):
        """
        examples: list of examples, each example is of form (board, pi, v)
        """
        input_boards, target_pis, target_vs = list(zip(*examples))
        input_boards = np.asarray(input_boards)
        target_pis = np.asarray(target_pis)
        target_vs = np.asarray(target_vs)
        self.nnet.model.fit(x=input_boards, y=[target_pis, target_vs], batch_size=args.batch_size, epochs=args.epochs, verbose=2)

    def predict(self, board):
        """
        board: np array with board
        """
        # timing
        #start = time.time()

        # preparing input
        board = board[np.newaxis, :, :]

        # run
        pi, v = self.nnet.model.predict(board)

        #print('PREDICTION TIME TAKEN : {0:03f}'.format(time.time()-start))
        return pi[0], v[0]

    def save_checkpoint(self, folder='checkpoint', filename='checkpoint.h5'):
        filepath = os.path.join(folder, filename)
        if not os.path.exists(folder):
            print(f"Checkpoint Directory does not exist! Making directory {folder}")
            os.mkdir(folder)
        else:
            print(f"Checkpoint Directory exists (at '{os.path.abspath(filepath)}').")
        self.nnet.model.save_weights(filepath)

    def load_checkpoint(self, folder='checkpoint', filename='checkpoint.h5'):
        filepath = os.path.join(folder, filename)
        if not os.path.exists(filepath):
            print(f"No model in path '{os.path.abspath(filepath)}'!")
            raise Exception("No model in path!")
        self.nnet.model.load_weights(filepath)

examples: list of examples, each example is of form (board, pi, v) board: np array with board # timing #start = time.time() # preparing input # run #print('PREDICTION TIME TAKEN : {0:03f}'.format(time.time()-start))