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 import layers
from keras import models
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 = 400

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)

    #flatten
    y_train = y_train.reshape(y_train.shape[0],image_wid * image_hig * image_dep)
    y_test = y_test.reshape(y_test.shape[0],image_wid * image_hig * image_dep)
    x_train = x_train.reshape(x_train.shape[0],image_wid * image_hig)
    x_test = x_test.reshape(x_test.shape[0], image_wid * image_hig)

    #正規化
    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
    
    model = models.Sequential()
    model.add(layers.Dense(900,activation='relu',
                               input_shape = (image_wid * image_hig,)))
    #model.add(layers.Dense(625,activation='relu'))
    model.add(layers.Dense(400,activation='relu'))
    #model.add(layers.Dense(625,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(x_train,y_train,
              epochs=epochs,
              batch_size=128,
              shuffle=True,
              )
    score = model.evaluate(x_test,y_test,verbose=0)
    print('test xentropy: ', score)

    z = model.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()