import pickle
import numpy as np
from matplotlib import pyplot as plt
import seaborn as sns
# def transform_boards(list_of_boards):
#     array_of_boards = np.array(list_of_boards)
#     one_hot_sparse = self.board_ohe.transform(array_of_boards.reshape([-1, 1]))
#     return one_hot_sparse.toarray().reshape([-1, 4, 4, 16])

def transform_boards1(list_of_boards):
    array_of_boards = np.array(list_of_boards)
    return array_of_boards.reshape([-1, 4, 4, 1])

def generate_rotated_boards(list_of_boards, y):
    boards = transform_boards1(list_of_boards)
    boards1 = np.rot90(boards, axes=(1, 2))
    boards2 = np.rot90(boards1, axes=(1, 2))
    boards3 = np.rot90(boards2, axes=(1, 2))
    flipped0 = np.flip(boards, axis=1)
    flipped1 = np.rot90(flipped0, axes=(1, 2))
    flipped2 = np.rot90(flipped1, axes=(1, 2))
    flipped3 = np.rot90(flipped2, axes=(1, 2))
    x_boards = np.concatenate([boards, boards1, boards2, boards3, flipped0, flipped1, flipped2, flipped3], axis=0)
    y_extended = np.array(y*8).reshape(-1, 1)
    return x_boards, y_extended
with open("game_history.pickle", 'rb') as ff:
    game_history = pickle.load(ff)
import keras as k
def conv_prelu(x, n=512, padding='valid'):
    x = k.layers.Conv2D(filters=n, kernel_size=(2, 2), padding=padding)(x)
    x = k.layers.PReLU()(x)
    return x



input_boards = k.layers.Input(shape=(4, 4, 1), name="input_boards")
x = conv_prelu(input_boards, padding='same')
x = conv_prelu(x, padding='same')
x = conv_prelu(x, padding='same')
x = conv_prelu(x) # 3, 3
x = conv_prelu(x) # 2, 2
x = conv_prelu(x) # 1, 1
x = k.layers.Flatten()(x)
x = k.layers.Dense(64, activation=k.layers.PReLU())(x)
x = k.layers.Dense(16, activation=k.layers.PReLU())(x)
prediction = k.layers.Dense(1, activation='linear')(x)
my_optimizer = k.optimizers.Adam(lr=0.001, beta_1=0.9, beta_2=0.999, epsilon=None, decay=0.001, amsgrad=False)
model = k.models.Model(inputs=input_boards, outputs=prediction)
model.compile(optimizer=my_optimizer, loss='mean_squared_error',metrics=['mae', 'logcosh'])

tb = k.callbacks.TensorBoard(log_dir="jp_logs/run1", write_graph=False)
boards = game_history[0]
values = game_history[1]
x, y = generate_rotated_boards(boards, values)
x.shape, y.shape
model.fit(x, y, verbose=1, epochs=30, batch_size=4000)