import cv2
from PIL import Image as Im
import numpy as np
import os
import glob
import keras
from keras.preprocessing.image import load_img
from keras.layers import Input,Conv2D,MaxPooling2D,UpSampling2D,Dense
from keras import layers
from keras import models
from keras.models import Model
from keras import callbacks
from keras.datasets import cifar100
from matplotlib import pyplot as plt
from matplotlib import cm

path_1 = './image2/*'
path_2 = './test_pic/*'
image_wid = 32
image_hig = 32
img_size = (image_wid, image_hig)
image_dep = 3

epochs = 100

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

def main():
    x_train = []
    x_test  = []
    y_train = []
    y_test  = []
    #パスとフォルダの取得
    image_list = glob.glob(path_1)
    image_list2 = glob.glob(path_2)

    #ファイルから画像ロード,教師データの作成
    for i in image_list:
        image = load_img(i,color_mode='rgb',target_size=img_size)
        image_array = np.asarray(image)
        y_train.append(image_array)
    y_train = np.array(y_train)
    #白黒入力データの作成
    for i in range(len(y_train)):
        x_train.append(graycale(y_train[i]))
    x_train = np.array(x_train)

    #テストデータの作成、教師画像
    for i in image_list2:
        image = load_img(i,color_mode='rgb',target_size=(img_size))
        image_array = np.asarray(image)
        y_test.append(image_array)
    y_test = np.array(y_test)
    #白黒画像の作成
    for i in range(len(y_test)):
        x_test.append(graycale(y_test[i]))
    x_test = np.array(x_test)

    #正規化
    x_train = x_train.astype('float32')
    x_test  =  x_test.astype('float32')
    x_train = x_train/256
    x_test =  x_test /256
    y_train = y_train.astype('float32')
    y_test  =  y_test.astype('float32')
    y_train = y_train/256
    y_test  = y_test /256

    x_train = np.reshape(x_train,(len(x_train),image_wid,image_hig,1))
    x_test = np.reshape(x_test,(len(x_test),image_wid,image_hig,1))
    y_train = np.reshape(y_train,(len(y_train),image_wid,image_hig,image_dep))
    y_test = np.reshape(y_test,(len(y_test),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)

    autoencoder.compile(optimizer='adam',loss='binary_crossentropy')
    history = autoencoder.fit(x_train,y_train,
              epochs=epochs,
              batch_size=128,
              shuffle=True,
              )
    score = autoencoder.evaluate(x_test,y_test,verbose=0)
    print('test xentropy: ', score)

    z = autoencoder.predict(x_test,verbose=0)
    x = x_test.reshape(x_test.shape[0],image_wid,image_hig)
    y = y_test.reshape(y_test.shape[0],image_wid,image_hig,image_dep)
    z = z.reshape(z.shape[0],image_wid,image_hig,image_dep)

    fig, axes = plt.subplots(3,10)
    for i in range(10):
        axes[0][i].imshow(x[i],cmap='gray')
        axes[0][i].axis('off')
        axes[1][i].imshow(y[i])
        axes[1][i].axis('off')
        axes[2][i].imshow(z[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()