# Path: betago/processor.py
# class SevenPlaneProcessor(GoDataProcessor):
#     '''
#     Implementation of a Go data processor, using seven planes of 19x19 values to represent the position of
#     a go board, as explained below.
# 
#     This closely reflects the representation suggested in Clark, Storkey:
#     http://arxiv.org/abs/1412.3409
#     '''
# 
#     def __init__(self, data_directory='data', num_planes=7, consolidate=True, use_generator=False):
#         super(SevenPlaneProcessor, self).__init__(data_directory=data_directory,
#                                                   num_planes=num_planes,
#                                                   consolidate=consolidate,
#                                                   use_generator=use_generator)
# 
#     def feature_and_label(self, color, move, go_board, num_planes):
#         '''
#         Parameters
#         ----------
#         color: color of the next person to move
#         move: move they decided to make
#         go_board: represents the state of the board before they moved
# 
#         Planes we write:
#         0: our stones with 1 liberty
#         1: our stones with 2 liberty
#         2: our stones with 3 or more liberties
#         3: their stones with 1 liberty
#         4: their stones with 2 liberty
#         5: their stones with 3 or more liberties
#         6: simple ko
#         '''
#         row, col = move
#         enemy_color = go_board.other_color(color)
#         label = row * 19 + col
#         move_array = np.zeros((num_planes, go_board.board_size, go_board.board_size))
#         for row in range(0, go_board.board_size):
#             for col in range(0, go_board.board_size):
#                 pos = (row, col)
#                 if go_board.board.get(pos) == color:
#                     if go_board.go_strings[pos].liberties.size() == 1:
#                         move_array[0, row, col] = 1
#                     elif go_board.go_strings[pos].liberties.size() == 2:
#                         move_array[1, row, col] = 1
#                     elif go_board.go_strings[pos].liberties.size() >= 3:
#                         move_array[2, row, col] = 1
#                 if go_board.board.get(pos) == enemy_color:
#                     if go_board.go_strings[pos].liberties.size() == 1:
#                         move_array[3, row, col] = 1
#                     elif go_board.go_strings[pos].liberties.size() == 2:
#                         move_array[4, row, col] = 1
#                     elif go_board.go_strings[pos].liberties.size() >= 3:
#                         move_array[5, row, col] = 1
#                 if go_board.is_simple_ko(color, pos):
#                     move_array[6, row, col] = 1
#         return move_array, label


from keras.models import Sequential
from keras.layers.core import Dense, Dropout, Activation, Flatten
from keras.layers.convolutional import Conv2D, MaxPooling2D
from keras.utils import np_utils
from betago.processor import SevenPlaneProcessor
nb_conv = 3  # convolution kernel size

# SevenPlaneProcessor loads seven planes (doh!) of 19*19 data points, so we need 7 input channels
processor = SevenPlaneProcessor()
input_channels = processor.num_planes

# Load go data and one-hot encode labels
X, y = processor.load_go_data(num_samples=1000)
X = X.astype('float32')
Y = np_utils.to_categorical(y, nb_classes)

# Specify a keras model with two convolutional layers and two dense layers,
# connecting the (num_samples, 7, 19, 19) input to the 19*19 output vector.
model = Sequential()
model.add(Conv2D(nb_filters, (nb_conv, nb_conv), padding='valid',
                 input_shape=(input_channels, go_board_rows, go_board_cols),
                 data_format='channels_first'))
model.add(Activation('relu'))
model.add(Conv2D(nb_filters, (nb_conv, nb_conv), data_format='channels_first'))
model.add(Activation('relu'))
model.add(MaxPooling2D(pool_size=(nb_pool, nb_pool)))
model.add(Dropout(0.25))
model.add(Flatten())
model.add(Dense(256))
model.add(Activation('relu'))
model.add(Dropout(0.5))
model.add(Dense(nb_classes))
model.add(Activation('softmax'))
model.compile(loss='categorical_crossentropy', optimizer='adadelta', metrics=['accuracy'])


# Path: betago/processor.py
# class SevenPlaneProcessor(GoDataProcessor):
#     '''
#     Implementation of a Go data processor, using seven planes of 19x19 values to represent the position of
#     a go board, as explained below.
# 
#     This closely reflects the representation suggested in Clark, Storkey:
#     http://arxiv.org/abs/1412.3409
#     '''
# 
#     def __init__(self, data_directory='data', num_planes=7, consolidate=True, use_generator=False):
#         super(SevenPlaneProcessor, self).__init__(data_directory=data_directory,
#                                                   num_planes=num_planes,
#                                                   consolidate=consolidate,
#                                                   use_generator=use_generator)
# 
#     def feature_and_label(self, color, move, go_board, num_planes):
#         '''
#         Parameters
#         ----------
#         color: color of the next person to move
#         move: move they decided to make
#         go_board: represents the state of the board before they moved
# 
#         Planes we write:
#         0: our stones with 1 liberty
#         1: our stones with 2 liberty
#         2: our stones with 3 or more liberties
#         3: their stones with 1 liberty
#         4: their stones with 2 liberty
#         5: their stones with 3 or more liberties
#         6: simple ko
#         '''
#         row, col = move
#         enemy_color = go_board.other_color(color)
#         label = row * 19 + col
#         move_array = np.zeros((num_planes, go_board.board_size, go_board.board_size))
#         for row in range(0, go_board.board_size):
#             for col in range(0, go_board.board_size):
#                 pos = (row, col)
#                 if go_board.board.get(pos) == color:
#                     if go_board.go_strings[pos].liberties.size() == 1:
#                         move_array[0, row, col] = 1
#                     elif go_board.go_strings[pos].liberties.size() == 2:
#                         move_array[1, row, col] = 1
#                     elif go_board.go_strings[pos].liberties.size() >= 3:
#                         move_array[2, row, col] = 1
#                 if go_board.board.get(pos) == enemy_color:
#                     if go_board.go_strings[pos].liberties.size() == 1:
#                         move_array[3, row, col] = 1
#                     elif go_board.go_strings[pos].liberties.size() == 2:
#                         move_array[4, row, col] = 1
#                     elif go_board.go_strings[pos].liberties.size() >= 3:
#                         move_array[5, row, col] = 1
#                 if go_board.is_simple_ko(color, pos):
#                     move_array[6, row, col] = 1
#         return move_array, label


# Path: examples/train_and_store.py
from keras.models import Sequential
from keras.layers.core import Dense, Dropout, Activation, Flatten
from keras.layers.convolutional import Conv2D, MaxPooling2D
from keras.utils import np_utils
from betago.processor import SevenPlaneProcessor

nb_conv = 3  # convolution kernel size

# SevenPlaneProcessor loads seven planes (doh!) of 19*19 data points, so we need 7 input channels
processor = SevenPlaneProcessor()
input_channels = processor.num_planes

# Load go data and one-hot encode labels
X, y = processor.load_go_data(num_samples=1000)
X = X.astype('float32')
Y = np_utils.to_categorical(y, nb_classes)

# Specify a keras model with two convolutional layers and two dense layers,
# connecting the (num_samples, 7, 19, 19) input to the 19*19 output vector.
model = Sequential()
model.add(Conv2D(nb_filters, (nb_conv, nb_conv), padding='valid',
                 input_shape=(input_channels, go_board_rows, go_board_cols),
                 data_format='channels_first'))
model.add(Activation('relu'))
model.add(Conv2D(nb_filters, (nb_conv, nb_conv), data_format='channels_first'))
model.add(Activation('relu'))
model.add(MaxPooling2D(pool_size=(nb_pool, nb_pool)))
model.add(Dropout(0.25))
model.add(Flatten())
model.add(Dense(256))
model.add(Activation('relu'))
model.add(Dropout(0.5))
model.add(Dense(nb_classes))
model.add(Activation('softmax'))
model.compile(loss='categorical_crossentropy', optimizer='adadelta', metrics=['accuracy'])


model.fit(X, Y, batch_size=batch_size, epochs=nb_epoch, verbose=1)