import numpy as np
from urllib import request
import gzip
import pickle

filename = [
["training_images","train-images-idx3-ubyte.gz"],
["test_images","t10k-images-idx3-ubyte.gz"],
["training_labels","train-labels-idx1-ubyte.gz"],
["test_labels","t10k-labels-idx1-ubyte.gz"]
]

def download_mnist():
    base_url = "http://yann.lecun.com/exdb/mnist/"
    for name in filename:
        print("Downloading "+name[1]+"...")
        request.urlretrieve(base_url+name[1], name[1])
    print("Download complete.")

def save_mnist():
    mnist = {}
    for name in filename[:2]:
        with gzip.open(name[1], 'rb') as f:
            mnist[name[0]] = np.frombuffer(f.read(), np.uint8, offset=16).reshape(-1,28*28)
    for name in filename[-2:]:
        with gzip.open(name[1], 'rb') as f:
            mnist[name[0]] = np.frombuffer(f.read(), np.uint8, offset=8)
    with open("mnist.pkl", 'wb') as f:
        pickle.dump(mnist,f)
    print("Save complete.")

def init_mnist():
    download_mnist()
    save_mnist()

def load_mnist():
    with open("mnist.pkl",'rb') as f:
        mnist = pickle.load(f)
    return mnist["training_images"], mnist["training_labels"], mnist["test_images"], mnist["test_labels"]


init_mnist()
x_train, t_train, x_test, t_test = load_mnist()
x_train = x_train.reshape(-1, 28, 28)
x_test = x_test.reshape(-1, 28, 28)

idx = 0  # <== SET AN INDEX
img = x_train[idx,:] # First image in the training set.
plt.imshow(img,cmap='gray')
plt.show() # Show the image

np.random.seed(42)  # <== Changing this value would generate different selections

# Pre-process the train set
x_train = x_train / np.max(x_train)    # normalization
x_train[x_train>=0.5] = 1              # apply threshold for binarization
x_train[x_train<0.5] = -1              # apply threshold for binarization

# Shuffle the train set
shuffle_mask = np.arange( len(t_train) )
np.random.shuffle(shuffle_mask)
x_train = x_train[shuffle_mask]
t_train = t_train[shuffle_mask]

# Select 10 different numbers from the train set
uniq, uniq_idx = np.unique(t_train, return_index=True)
numbers = x_train[uniq_idx]

f, axs = plt.subplots(2, 5)
for i, ax in enumerate(axs.flatten()):
    ax.imshow(numbers[i], cmap="gray")
    ax.set_title("Value = "+ str(i))
f.tight_layout()

# Add noise to samples
noise = 100
noisy_numbers = numbers.copy()
for number in noisy_numbers:
    mask = np.ones((28,28))
    for i in range(noise):
        r,c = np.random.randint(0,28, size=2)
        number[r,c] *= -1

f, axs = plt.subplots(2, 5)
for i, ax in enumerate(axs.flatten()):
    ax.imshow(noisy_numbers[i], cmap="gray")
    ax.set_title("Value = "+ str(i))
f.tight_layout()

class HopfieldNetwork:
    def __init__(self, k, threshold=0):
        self.k = k
        self.threshold = threshold

    def calculate_energy(self, x):
        return -0.5*np.sum(np.matmul(x.T, x)*self.w) + np.sum(x*self.threshold)

    def train(self, data):
        self.w = np.zeros((self.k, self.k))
        for x in data:
            self.w += self.calculate_w_per_sample(x)
        #print("w:", self.w)

    def activation(self, z):
        if z > self.threshold:
            return 1
        return -1

    def predict(self, x):   #update
        x = x.copy()
        x = x.reshape(1, -1)

        #async update
        max_epoch = 1
        counter = 0 # counter or not changed states
        for e in range(max_epoch):
            if (counter > self.k):
                break

            order = np.arange(self.k)
            np.random.shuffle(order)

            for i in order:
                x_i = x @ self.w[i]
                x_i_a = self.activation(x_i)
                
                if (x_i_a != x[0,i]):
                    counter += 1
                    x[0,i] = x_i_a
                    if (self.k < 101):
                        print("State:", x)
                    else:
                        print("Energy:", self.calculate_energy(x))

        return x
        
    def calculate_w_per_sample(self, x):
        x = x.reshape(1, -1)
        n = x.shape[1]
        w = np.matmul(x.T, x) * (1 - np.eye(n))
        return w

data = np.array([ [1,1,1,1,-1,-1,-1], [-1,1,-1,1,-1,1,-1], [-1,-1,1,1,-1,1,-1] ])
print("Printing only changed states...")
hn = HopfieldNetwork(7)
hn.train(data)
final_state = hn.predict(np.array([-1,-1,-1,1,-1,-1,-1]))
print("final state:", final_state)

hn = HopfieldNetwork(28*28)
hn.train(numbers[[0,1]])

f, axs = plt.subplots(1, 2)
axs[0].imshow(numbers[0], cmap="gray")
axs[0].set_title("Train Pattern 0")
axs[1].imshow(numbers[1], cmap="gray")
axs[1].set_title("Train Pattern 1")
f.tight_layout()

f, axs = plt.subplots(1, 2)

out = hn.predict(noisy_numbers[0]).reshape(28, 28)
axs[0].imshow(noisy_numbers[0], cmap="gray")
axs[0].set_title("Input")
axs[1].imshow(out, cmap="gray")
axs[1].set_title("Output")

f.tight_layout()

f, axs = plt.subplots(1, 2)

out = hn.predict(noisy_numbers[1]).reshape(28, 28)
axs[0].imshow(noisy_numbers[1], cmap="gray")
axs[0].set_title("Input")
axs[1].imshow(out, cmap="gray")
axs[1].set_title("Output")

f.tight_layout()

sign_T = np.array([[1,1,1],[-1,1,-1],[-1,1,-1]])
sign_plus = np.array([[-1,1,-1], [1,1,1], [-1,1,-1]])
sign_x = np.array([[1,-1,1],[-1,1,-1],[1,-1,1]])
data = np.array([sign_T ,sign_plus, sign_x]).reshape(-1,9)

hn = HopfieldNetwork(9)
hn.train(data)

f, axs = plt.subplots(1, 3)

axs[0].imshow(sign_T, cmap="gray")
axs[0].set_title("data[0]")
axs[1].imshow(sign_plus, cmap="gray")
axs[1].set_title("data[1]")
axs[2].imshow(sign_x, cmap="gray")
axs[2].set_title("data[2]")

f.tight_layout()

# TEST-1
test_data = np.array([[1,-1,1],[-1,-1,-1],[-1,-1,1]])
final_state = hn.predict(test_data)
print("final state:", final_state)

f, axs = plt.subplots(1, 2)
axs[0].imshow(test_data, cmap="gray")
axs[0].set_title("Input")
axs[1].imshow(final_state.reshape(3,3), cmap="gray")
axs[1].set_title("Output")
f.tight_layout()

# TEST-2
test_data = np.array([[1,1,1],[1,1,1],[1,1,1]])
final_state = hn.predict(test_data)
print("final state:", final_state)

f, axs = plt.subplots(1, 2)
axs[0].imshow(test_data, cmap="gray")
axs[0].set_title("Input")
axs[1].imshow(final_state.reshape(3,3), cmap="gray")
axs[1].set_title("Output")
f.tight_layout()

# TEST-2
test_data = np.array([[1,1,1],[-1,1,1],[-1,-1,-1]])
final_state = hn.predict(test_data)
print("final state:", final_state)

f, axs = plt.subplots(1, 2)
axs[0].imshow(test_data, cmap="gray")
axs[0].set_title("Input")
axs[1].imshow(final_state.reshape(3,3), cmap="gray")
axs[1].set_title("Output")
f.tight_layout()