## Author: SHIVANK SHARMA
## Modified: Toshio Iwata
## Released under Apache 2.0 license 

import numpy as np
import pandas as pd
import os
import cv2
import matplotlib.pyplot as plt
import tensorflow as tf
from tensorflow import keras
from PIL import Image
from sklearn.model_selection import train_test_split
from tensorflow.keras.preprocessing.image import ImageDataGenerator
from tensorflow.keras.optimizers import Adam
from sklearn.metrics import accuracy_score
np.random.seed(42)

from matplotlib import style
style.use('fivethirtyeight')

## Modified paths, be careful of the location of Train and Test folder
data_dir = './'
train_path = './train_1'
test_path = './'

# Resizing the images to 30x30x3
IMG_HEIGHT = 30
IMG_WIDTH = 30
channels = 3

NUM_CATEGORIES = len(os.listdir(train_path))
NUM_CATEGORIES

# Label Overview
### Modified for lisa
classes = { 0:'zero',
            1:'one',
            2:'two', 
            3:'three', 
            4:'four', 
            5:'five', 
            6:'six', 
            7:'seven', 
            8:'eight',
            9:'nine',
            10:'white', 
            11:'black', 
            12:'wh2bk', 
            13:'bk2wh', 
            14:'wh_bk', 
            15:'bk_wh' }


folders = os.listdir(train_path)

train_number = []
class_num = []

for folder in folders:
    train_files = os.listdir(train_path + '/' + folder)
    train_number.append(len(train_files))
    class_num.append(classes[int(folder)])
    
# Sorting the dataset on the basis of number of images in each class
zipped_lists = zip(train_number, class_num)
sorted_pairs = sorted(zipped_lists)
tuples = zip(*sorted_pairs)
train_number, class_num = [ list(tuple) for tuple in  tuples]

# Plotting the number of images in each class
plt.figure(figsize=(21,10))  
plt.bar(class_num, train_number)
plt.xticks(class_num, rotation='vertical')
plt.show()

# Visualizing 25 random images from test data
import random
from matplotlib.image import imread

'''
imgs = 'ds1_pic29.jpg' ### test["Path"].values
plt.figure(figsize=(25,25))
##plt.subplot(5,5,i)

img_path = train_path + '/0/' + imgs
rand_img = imread(img_path)
plt.imshow(rand_img)
#plt.grid(b=None)
plt.xlabel(rand_img.shape[1], fontsize = 20)#width of image


test = pd.read_csv(data_dir + '/Test.csv')
imgs = test["Path"].values

plt.figure(figsize=(25,25))
for i in range(1,26):
    plt.subplot(5,5,i)
    random_img_path = data_dir + '/' + random.choice(imgs)
    rand_img = imread(random_img_path)
    plt.imshow(rand_img)
    plt.grid(b=None)
    plt.xlabel(rand_img.shape[1], fontsize = 20)#width of image
   
'''

## Modified for Shiroume
def resize (image, size):
    h, w = image.shape[:2]
    aspect = w / h
    nh = nw = size
    if 1 >= aspect:
        nw = round(nh * aspect)
    else:
        nh = round(nw / aspect)
    resized = cv2.resize(image, dsize=(nw, nh))
    return resized

## Modified for Shiroume
def shiroume (image, size):
    resized = resize(image, size)

    h, w = resized.shape[:2]
    x = y = 0
    if h < w:
        y = (size - h) // 2
    else:
        x = (size - w) // 2
    resized = Image.fromarray(resized) #一旦PIL形式に変換
    canvas = Image.new(resized.mode, (size, size), (255, 255, 255)) # 黒埋めなら(0,0,0)にする
    canvas.paste(resized, (x, y))
    dst = np.array(canvas) #numpy(OpenCV)形式に戻す
    return dst


image_data = []
image_labels = []

for i in range(NUM_CATEGORIES):
    path = data_dir + '/train_1/' + str(i)
    images = os.listdir(path)

    for img in images:
        try:
            image = cv2.imread(path + '/' + img)

            ## Modified for Shiroume
            h, w = image.shape[:2]
            if h < w :
                val = w
            else :
                val = h
            ume_image = shiroume(image, val) #リサイズしたいサイズを指定

            image_fromarray = Image.fromarray(image, 'RGB') ## ume_image for shiroume
            resize_image = image_fromarray.resize((IMG_HEIGHT, IMG_WIDTH))
            image_data.append(np.array(resize_image))
            image_labels.append(i)
        except:
            print("Error in " + img)

# Changing the list to numpy array
image_data = np.array(image_data)
image_labels = np.array(image_labels)

print(image_data.shape, image_labels.shape)

shuffle_indexes = np.arange(image_data.shape[0])
np.random.shuffle(shuffle_indexes)
image_data = image_data[shuffle_indexes]
image_labels = image_labels[shuffle_indexes]

X_train, X_val, y_train, y_val = train_test_split(image_data, image_labels, test_size=0.3, random_state=42, shuffle=True)

X_train = X_train/255 
X_val = X_val/255

print("X_train.shape", X_train.shape)
print("X_valid.shape", X_val.shape)
print("y_train.shape", y_train.shape)
print("y_valid.shape", y_val.shape)

y_train = keras.utils.to_categorical(y_train, NUM_CATEGORIES)
y_val = keras.utils.to_categorical(y_val, NUM_CATEGORIES)

print(y_train.shape)
print(y_val.shape)

model = keras.models.Sequential([    
    keras.layers.Conv2D(filters=16, kernel_size=(3,3), activation='relu', input_shape=(IMG_HEIGHT,IMG_WIDTH,channels)),
    keras.layers.Conv2D(filters=32, kernel_size=(3,3), activation='relu'),
    keras.layers.MaxPool2D(pool_size=(2, 2)),
    ##keras.layers.BatchNormalization(axis=-1),
    
    keras.layers.Conv2D(filters=64, kernel_size=(3,3), activation='relu'),
    keras.layers.Conv2D(filters=128, kernel_size=(3,3), activation='relu'),
    keras.layers.MaxPool2D(pool_size=(2, 2)),
    ##keras.layers.BatchNormalization(axis=-1),
    
    keras.layers.Flatten(),
    ##keras.layers.Dense(512, activation='relu'),
    ##keras.layers.Dense(256, activation='relu'),
    ## Modified from 512 to 128
    keras.layers.Dense(128, activation='relu'),
    ##keras.layers.BatchNormalization(),
    keras.layers.Dropout(rate=0.5),
    
    ##keras.layers.Dense(43, activation='softmax')
    keras.layers.Dense(16, activation='softmax') ### Modified for lisa
])

lr = 0.001
epochs = 30

####opt = Adam(lr=lr, decay=lr / (epochs * 0.5)) # for older jupyter lab
opt = tf.keras.optimizers.legacy.Adam(lr=lr, decay=lr / (epochs * 0.5)) # for newer jupyter lab

model.compile(loss='categorical_crossentropy', optimizer=opt, metrics=['accuracy'])

aug = ImageDataGenerator(
    rotation_range=10,
    zoom_range=0.15,
    width_shift_range=0.1,
    height_shift_range=0.1,
    shear_range=0.15,
    horizontal_flip=False,
    vertical_flip=False,
    fill_mode="nearest")

history = model.fit(aug.flow(X_train, y_train, batch_size=32), epochs=epochs, validation_data=(X_val, y_val))

############################### by iwata start

import csv

###################################################
## for write
###################################################
########################### lay1
lay1 = model.layers[0]
print("lay1 weight shape : ", lay1.get_weights()[0].shape)
print("lay1 weight shape : ", lay1.get_weights()[1].shape)

with open("w1.csv", mode='w') as f:
    for i in range(0, lay1.get_weights()[0].shape[2]):
        for r in range(0, lay1.get_weights()[0].shape[0]):
            for c in range(0, lay1.get_weights()[0].shape[1]):
                for k in range(0, lay1.get_weights()[0].shape[3]):
                    tmpval = lay1.get_weights()[0][r][c][i][k] ### [0]に重み係数
                    f.write(str(tmpval))
                    f.write(', ')
                f.write('\n')

with open("b1.csv", mode='w') as f:
    for c in range(0, lay1.get_weights()[1].shape[0]):
        tmpval = lay1.get_weights()[1][c] ### [1]にバイアス
        f.write(str(tmpval))
        f.write(', ')

########################### lay2
lay2 = model.layers[1]
print("lay2 weight shape : ", lay2.get_weights()[0].shape)
print("lay2 weight shape : ", lay2.get_weights()[1].shape)

with open("w2.csv", mode='w') as f:
    for i in range(0, lay2.get_weights()[0].shape[2]):
        for r in range(0, lay2.get_weights()[0].shape[0]):
            for c in range(0, lay2.get_weights()[0].shape[1]):
                for k in range(0, lay2.get_weights()[0].shape[3]):
                    tmpval = lay2.get_weights()[0][r][c][i][k] ### [0]に重み係数
                    f.write(str(tmpval))
                    f.write(', ')
                f.write('\n')

with open("b2.csv", mode='w') as f:
    for c in range(0, lay2.get_weights()[1].shape[0]):
        tmpval = lay2.get_weights()[1][c] ### [1]にバイアス
        f.write(str(tmpval))
        f.write(', ')

########################### lay3
lay3 = model.layers[3]
print("lay3 weight shape : ", lay3.get_weights()[0].shape)
print("lay3 weight shape : ", lay3.get_weights()[1].shape)

with open("w3.csv", mode='w') as f:
    for i in range(0, lay3.get_weights()[0].shape[2]):
        for r in range(0, lay3.get_weights()[0].shape[0]):
            for c in range(0, lay3.get_weights()[0].shape[1]):
                for k in range(0, lay3.get_weights()[0].shape[3]):
                    tmpval = lay3.get_weights()[0][r][c][i][k] ### [0]に重み係数
                    f.write(str(tmpval))
                    f.write(', ')
                f.write('\n')

with open("b3.csv", mode='w') as f:
    for c in range(0, lay3.get_weights()[1].shape[0]):
        tmpval = lay3.get_weights()[1][c] ### [1]にバイアス
        f.write(str(tmpval))
        f.write(', ')

########################### lay4
lay4 = model.layers[4]
print("lay4 weight shape : ", lay4.get_weights()[0].shape)
print("lay4 weight shape : ", lay4.get_weights()[1].shape)

with open("w4.csv", mode='w') as f:
    for i in range(0, lay4.get_weights()[0].shape[2]):
        for r in range(0, lay4.get_weights()[0].shape[0]):
            for c in range(0, lay4.get_weights()[0].shape[1]):
                for k in range(0, lay4.get_weights()[0].shape[3]):
                    tmpval = lay4.get_weights()[0][r][c][i][k] ### [0]に重み係数
                    f.write(str(tmpval))
                    f.write(', ')
                f.write('\n')

with open("b4.csv", mode='w') as f:
    for c in range(0, lay4.get_weights()[1].shape[0]):
        tmpval = lay4.get_weights()[1][c] ### [1]にバイアス
        f.write(str(tmpval))
        f.write(', ')

########################### lay5
lay5 = model.layers[7]
print("lay5 weight shape : ", lay5.get_weights()[0].shape)
print("lay5 weight shape : ", lay5.get_weights()[1].shape)

with open("w5.csv", mode='w') as f:
    for r in range(0, lay5.get_weights()[0].shape[0]):
        for c in range(0, lay5.get_weights()[0].shape[1]):
            tmpval = lay5.get_weights()[0][r][c] ### [0]に重み係数
            f.write(str(tmpval))
            f.write(', ')
        f.write('\n')

with open("b5.csv", mode='w') as f:
    for c in range(0, lay5.get_weights()[1].shape[0]):
        tmpval = lay5.get_weights()[1][c] ### [1]にバイアス
        f.write(str(tmpval))
        f.write(', ')

########################### lay6
lay6 = model.layers[9]
print("lay6 weight shape : ", lay6.get_weights()[0].shape)
print("lay6 weight shape : ", lay6.get_weights()[1].shape)

with open("w6.csv", mode='w') as f:
    for r in range(0, lay6.get_weights()[0].shape[0]):
        for c in range(0, lay6.get_weights()[0].shape[1]):
            tmpval = lay6.get_weights()[0][r][c] ### [0]に重み係数
            f.write(str(tmpval))
            f.write(', ')
        f.write('\n')

with open("b6.csv", mode='w') as f:
    for c in range(0, lay6.get_weights()[1].shape[0]):
        tmpval = lay6.get_weights()[1][c] ### [1]にバイアス
        f.write(str(tmpval))
        f.write(', ')

############################### by iwata end


pd.DataFrame(history.history).plot(figsize=(8, 5))
plt.grid(True)
plt.gca().set_ylim(0, 1)
plt.show()

test = pd.read_csv(data_dir + '/TestMnist.csv') ### Modified for lisa

labels = test["ClassId"].values
imgs = test["Path"].values

data =[]

for img in imgs:
    try:
        image = cv2.imread(data_dir + '/' +img)

        ## Modified for Shiroume
        h, w = image.shape[:2]
        if h < w :
            val = w
        else :
            val = h
        ume_image = shiroume(image, val) #リサイズしたいサイズを指定

        image_fromarray = Image.fromarray(image, 'RGB') ## ume_image for shiroume
        resize_image = image_fromarray.resize((IMG_HEIGHT, IMG_WIDTH))
        data.append(np.array(resize_image))
    except:
        print("Error in " + img)
X_test = np.array(data)
X_test = X_test/255

####pred = model.predict_classes(X_test)
pred = np.argmax(model.predict(X_test), axis=-1) # for newer tensorflow

#Accuracy with the test data
print('Test Data accuracy: ',accuracy_score(labels, pred)*100)

from sklearn.metrics import confusion_matrix
cf = confusion_matrix(labels, pred)

import seaborn as sns
df_cm = pd.DataFrame(cf, index = classes,  columns = classes)
##plt.figure(figsize = (20,20))
sns.heatmap(df_cm, annot=True)

from sklearn.metrics import classification_report

print(classification_report(labels, pred))

plt.figure(figsize = (25, 25))

### Modified for lisa 
### changed index of images to show up
for i in range(48):
    plt.subplot(6, 8, i + 1)
    plt.grid(False)
    plt.xticks([])
    plt.yticks([])
    plt.rcParams["font.size"] = 16
    if i%16 == 0 :
        rndval = int(i/16)
    elif i%16 == 1 :
        rndval = 980 + int(i/16)
    elif i%16 == 2 :
        rndval = 2115 + 0 + int(i/16)
    elif i%16 == 3 :
        rndval = 3147 + 3 + int(i/16)
    elif i%16 == 4 :
        rndval = 4157 + 6 + int(i/16)
    elif i%16 == 5 :
        rndval = 5139 + 9 + int(i/16)
    elif i%16 == 6 :
        rndval = 6031 + 12 + int(i/16)
    elif i%16 == 7 :
        rndval = 6989 + 15 + int(i/16)
    elif i%16 == 8 :
        rndval = 8017 + 18 + int(i/16)
    elif i%16 == 9 :
        rndval = 8991 + int(i/16)
    elif i%16 == 10 :
        rndval = 10000 + int(i/16)
    elif i%16 == 11 :
        rndval = 10003 + int(i/16)
    elif i%16 == 12 :
        rndval = 10006 + int(i/16)
    elif i%16 == 13 :
        rndval = 10009 + int(i/16)
    elif i%16 == 14 :
        rndval = 10012 + int(i/16)
    elif i%16 == 15 :
        rndval = 10015 + int(i/16)
    prediction = pred[rndval]
    actual = labels[rndval]
    col = 'g'
    if prediction != actual:
        col = 'r'
    plt.xlabel('Actual={} || Pred={}'.format(actual, prediction), color = col)
    for m in range(0, X_test[rndval].shape[0]): 
        for n in range(0, X_test[rndval].shape[1]): 
            tmpg = X_test[rndval][m][n][0]
            tmpr = X_test[rndval][m][n][2]
            X_test[rndval][m][n][0] = tmpr
            X_test[rndval][m][n][2] = tmpg
    plt.imshow(X_test[rndval])
plt.show()