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

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

epochs = 400

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 graycale(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')

    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.array(p_train)

    xtest = []
    ytest = []
    img_list = glob.glob(path)
    for i in img_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(graycale(ytest[i]))
    xtest = np.array(xtest)

    p_train = p_train.reshape(p_train.shape[0],image_wid * image_hig * image_dep)
    z_train = z_train.reshape(z_train.shape[0],image_wid * image_hig)

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

    ytest = ytest.reshape(ytest.shape[0],image_wid * image_hig * image_dep)
    xtest = xtest.reshape(xtest.shape[0],image_wid * image_hig)
    xtest = xtest.astype('float32')
    ytest = ytest.astype('float32')
    xtest = xtest/256
    ytest = ytest/256

    model = models.Sequential()
    model.add(layers.Dense(900,activation='relu',
                               input_shape = (image_wid * image_hig,)))
    model.add(layers.Dense(400,activation='relu'))
    model.add(layers.Dense(900,activation='relu'))
    model.add(layers.Dense(image_wid * image_hig * image_dep,
                               activation='sigmoid'))
    model.compile(optimizer='adam',loss='binary_crossentropy')

    history = model.fit(z_train,p_train,
              epochs=epochs,
              batch_size=128,
              shuffle=True,
              )

    score = model.evaluate(xtest,ytest,verbose=0)
    print('test xentropy: ', score)

    z_test_sampled = xtest
    p_test_sampled = ytest

    z_test_sampled_pred = model.predict(z_test_sampled,verbose=0)
    z_test_sampled = z_test_sampled.reshape(z_test_sampled.shape[0],image_wid,image_hig)
    p_test_sampled = p_test_sampled.reshape(p_test_sampled.shape[0],image_wid,image_hig,image_dep)
    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()