import cv2
import numpy as np
import glob
from keras import layers
from keras import models
from keras import callbacks
from keras.preprocessing.image import load_img
from keras.datasets import cifar100
from matplotlib import pyplot as plt
from matplotlib import cm

from keras.layers import Input,Conv2D,MaxPooling2D,UpSampling2D,Dense
from keras.models import Model
from keras import backend as K

image_wid = 32
image_hig = 32
image_dep = 3
scale = 1
path = './test_pic/*'
img_size = (image_wid,image_hig)

epochs = 100

def resize(image):
    img = image.repeat(scale,axis=0).repeat(scale,axis=1)
    img = cv2.cvtColor(img,cv2.COLOR_BGR2GRAY)
    return np.asarray(img)

def resize_img(image):
    img = image.repeat(scale,axis=0).repeat(scale,axis=1)
    return np.asarray(img)

def grayscale(image):
    img = cv2.cvtColor(image,cv2.COLOR_BGR2GRAY)
    return np.asarray(img)

def main():
    (x_train,y_train),(x_test,y_test) = cifar100.load_data(label_mode='fine')

    test_list = glob.glob(path)

    z_train = []
    for i in range(len(x_train)):
        z_train.append(resize(x_train[i]))
    z_train = np.array(z_train)

    p_train=[]
    for i in range(len(x_train)):
        p_train.append(resize_img(x_train[i]))

    p_train = np.asarray(p_train)

    xtest = []
    ytest = []
    for i in test_list:
        image = load_img(i,color_mode='rgb',target_size=(img_size))
        image_array = np.asarray(image)
        ytest.append(image_array)
    ytest = np.array(ytest)

    for i in range(len(ytest)):
        xtest.append(grayscale(ytest[i]))
    xtest = np.array(xtest)

    z_train = z_train.astype('float32')
    z_train = z_train/256
    p_train = p_train.astype('float32')
    p_train = p_train/256

    xtest = xtest.astype('float32')
    xtest = xtest/256
    ytest = ytest.astype('float32')
    ytest = ytest/256

    z_train = np.reshape(z_train,(len(z_train),image_wid,image_hig,1))
    p_train = np.reshape(p_train,(len(p_train),image_wid,image_hig,image_dep))

    xtest = np.reshape(xtest,(len(xtest),image_wid,image_hig,1))
    ytest = np.reshape(ytest,(len(ytest),image_wid,image_hig,image_dep))

    input_img = Input(shape=(image_wid,image_hig,1))

    x = Conv2D(64,(3,3),activation='relu',padding='same')(input_img)
    x = MaxPooling2D((2,2),padding='same')(x)
    x = Conv2D(32,(3,3),activation='relu',padding='same')(x)
    encoded = MaxPooling2D((2,2),padding='same')(x)
    
    encoded = Dense(512,activation='relu')(encoded)

    x = Conv2D(32,(3,3),activation='relu',padding='same')(encoded)
    x = UpSampling2D((2,2))(x)
    x = Conv2D(64,(3,3),activation='relu',padding='same')(x)    
    x = UpSampling2D((2,2))(x)
    decoded = Conv2D(3,(3,3),activation='sigmoid',padding='same')(x)

    autoencoder = Model(input_img,decoded)

    print(autoencoder.summary)
    autoencoder.compile(optimizer='adam',loss='binary_crossentropy')
    history = autoencoder.fit(z_train,p_train,
              epochs=epochs,
              batch_size=128,
              shuffle=True,
              )

    score = autoencoder.evaluate(xtest,ytest,verbose=0)

    print('test xentropy: ', score)

    z_test_sampled = xtest
    p_test_sampled = ytest

    z_test_sampled_pred = autoencoder.predict(z_test_sampled,verbose=0)

    z_test_sampled = z_test_sampled.reshape(z_test_sampled.shape[0],image_wid,image_hig)
    z_test_sampled_pred = z_test_sampled_pred.reshape(z_test_sampled_pred.shape[0],image_wid,image_hig,image_dep)

    
    fig, axes = plt.subplots(3,10)
    for i in range(10):
        axes[0][i].imshow(z_test_sampled[i],cmap='gray')
        axes[0][i].axis('off')
        axes[1][i].imshow(p_test_sampled[i])
        axes[1][i].axis('off')
        axes[2][i].imshow(z_test_sampled_pred[i])
        axes[2][i].axis('off')
    plt.show()

    plt.plot(range(1,epochs+1),history.history['loss'],"o-",label="loss",)
    plt.ylim(ymin=0)
    plt.grid(color='gray')
    plt.title('model loss')
    plt.ylabel('Loss')
    plt.xlabel('Epoch')
    plt.legend(['Train'],loc = 'upper right')
    plt.show()

if __name__ == "__main__":
    main()