import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import tensorflow as tf
import random
#import image_dataset_loader


# fn.keras.datasets.fashion_mnist
mnist = tf.keras.datasets.mnist # mnist data of grayscale images of handwritten digits
#mnist = tf.keras.datasets.fashion_mnist
(x_train, y_train), (x_test, y_test) = mnist.load_data() # load it into training and testing sets


#mnist = tf.keras.datasets.fashion_mnist
#(_, _), (x_test, y_test) = mnist.load_data() # load it into training and testing sets

#(x_test, y_test), = image_dataset_loader.load('./data', ['train'],x_dtype='uint8', y_dtype='uint32')

#x_train = x_train[..., 0]
#x_test = x_train
# aca tiene que qreemplazar las variables por numpy creadas por mi

#x_train = x_train > 128
#x_test  = x_test  > 128



#print(y_train[:10])
x_train, x_test = x_train / 255.0, x_test/255.0 # the data is in the range (0,255) grayscale value for each pixel, normalize it
print(x_train.shape, x_test.shape) # look at how many images are there



#print(y_train.shape, y_test.shape)
#sprint(y_train[:10])

#for i in range(1):
#    plt.figure()
#    plt.imshow(x_train[i], cmap='gray')
#plt.show()



print(x_train.shape)

#---------------------------------------------------

modelAE = tf.keras.models.Sequential([
        # the input layer
        tf.keras.layers.InputLayer((28,28)),
        # flatten it into a single column
        tf.keras.layers.Flatten(), 
        # bottleneck layer
        tf.keras.layers.Dense(128, activation='relu'), #relu o sigmoid
        tf.keras.layers.Dense(64, activation='relu'), #relu o sigmoid
        tf.keras.layers.Dense(128, activation='relu'), #relu o sigmoid
        # the output layer
        #tf.keras.layers.Dense(1, activation='sigmoid'), #relu o sigmoid
        # the output layer, reshape the flattened column into an image shape
        #tf.keras.layers.Reshape((1, 1, 1))

        tf.keras.layers.Dense(28*28, activation='sigmoid'), #relu o sigmoid
        # the output layer, reshape the flattened column into an image shape
        tf.keras.layers.Reshape((28, 28, 1))

])



print(modelAE.summary())
model = modelAE
# compile, using the adamax optimizer for gradient descend and mean square error for loss
model.compile(optimizer='adamax', loss='mse')
# fit (train) the model on the test data. The input and the output is the same because
# we are trying to re-construct the original input after compression
#train_hist = model.fit(x_train[:3], np.array([0, .5, 1]) , epochs=5) # idea de autoencoder



train_hist = model.fit( x_train  > .5 ,   x_train , epochs=7)
x_test =  x_test > .5
print(train_hist)

#---------------------------------------------------

# arma la imagen

original = []
decoded = []
n=0
# predict the first 10 items from the test set
#w = random.randint(0, 10000)
w = 100
for img in x_test[ w:w+3 ]:
    #predict
    img1 = img
    pimg = model.predict(tf.expand_dims( img1, 0))
    #print(pimg)
    original.append(img1) # original
    decoded.append(pimg.squeeze()) # predicted
    n += 1

# ---- muestra --------
plt.figure(figsize=(12, 12))
for i in range(3):
    # display original image
    ax = plt.subplot(3, n, i + 1)
    ax.set_title("Original")
    plt.imshow(original[i])
    plt.gray()
    ax.get_xaxis().set_visible(False)
    ax.get_yaxis().set_visible(False)

    # display reconstruction
    ax = plt.subplot(3, n, i + 1 + n)
    ax.set_title("Decoded")
    plt.imshow(decoded[i])
    plt.gray()

    ax.get_xaxis().set_visible(False)
    ax.get_yaxis().set_visible(False)

    ax = plt.subplot(3, n, i + 1 + n + n)

    ax.set_title("Entrenado")
    plt.imshow(x_train[i])
    plt.gray()

    ax.get_xaxis().set_visible(False)
    ax.get_yaxis().set_visible(False)

plt.show()