import numpy as np
import matplotlib.pyplot as plt
from neupy import algorithms

def image(matrix):
    plt.imshow(matrix.reshape(5,5), cmap=plt.cm.binary);
    plt.axis('off')
    plt.show()

plus = np.matrix([
    0, 0, 1, 0, 0,
    0, 0, 1, 0, 0,
    1, 1, 1, 1, 1,
    0, 0, 1, 0, 0,
    0, 0, 1, 0, 0,
])

times = np.matrix([
    1, 0, 0, 0, 1,
    0, 1, 0, 1, 0,
    0, 0, 1, 0, 0,
    0, 1, 0, 1, 0,
    1, 0, 0, 0, 1,
])

minus = np.matrix([
    0, 0, 0, 0, 0,
    0, 0, 0, 0, 0,
    1, 1, 1, 1, 1,
    0, 0, 0, 0, 0,
    0, 0, 0, 0, 0,
])

obelus = np.matrix([
    0, 0, 1, 0, 0,
    0, 0, 0, 0, 0,
    1, 1, 1, 1, 1,
    0, 0, 0, 0, 0,
    0, 0, 1, 0, 0,
])

image(plus)
image(times)
image(minus)
image(obelus)

operations = np.concatenate([plus, times, minus, obelus])
nn = algorithms.DiscreteHopfieldNetwork()
nn.train(operations)

part = np.matrix([
    0, 0, 0, 0, 1,
    0, 0, 0, 0, 0,
    0, 0, 0, 0, 0,
    0, 1, 0, 0, 0,
    1, 0, 0, 0, 1,
])
image(part)

result = nn.predict(part)
image(result)