import numpy as np
import sklearn.model_selection as skl_ms
from keras.models import Sequential
from keras.layers import Conv2D
from keras.layers import Activation
from keras.layers.core import Reshape
from keras.utils import plot_model
from matplotlib import cm
import matplotlib.pyplot as plt

def plotArray(w):
    plt.imshow(w, cmap=cm.gray);
    plt.colorbar();
    plt.show();

k=8;
n=1000;
X = np.random.randint(0, k, [1000, 1, 2, 1]);
Y = np.zeros([1000, 1]);
epochs=50;

#Harder AND
for i in range(0, n):
    if X[i, 0, 0, 0] > 3 and X[i, 0, 1, 0] > 3:
        Y[i] = 1;

(X_train, X_test, Y_train, Y_test) = skl_ms.train_test_split(X, Y);

model = Sequential([Conv2D(1, kernel_size=[1,2], activation='sigmoid',     input_shape=[1, 2, 1]),
                           Reshape(target_shape=(1,))]);

plot_model(model, to_file='logic_model.png', show_shapes=True,     show_layer_names=True);

model.compile(optimizer='nadam', loss='mse');

print(model.summary());

model.fit(x=X_train, y=Y_train, epochs=epochs, batch_size=30);


#Generate a plot to see what the model predictions are as a function of input 1
#and input 2
A=np.zeros([k, k]);
for i in range(0, k):
    for j in range(0, k):
        b = np.array([i, j]);
        b = b.reshape([1, 1, 2, 1]);

        A[i, j] = model.predict(b);

plotArray(A);
plotArray(A>0.5);