import time

from numpy.random import seed
seed(8) #1

import tensorflow
tensorflow.random.set_seed(7)
# tensorflow.random.set_random_seed(7)

import numpy as np # linear algebra
import pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)

import os

from tensorflow.keras import backend as K
from tensorflow.keras.models import Model ,load_model
from tensorflow.keras.layers import Flatten, Dense, Dropout
from tensorflow.keras.applications.inception_resnet_v2 import InceptionResNetV2, preprocess_input
from keras.applications.vgg16 import preprocess_input
from keras.applications.vgg16 import decode_predictions
from keras.applications.vgg16 import VGG16
from tensorflow.keras.optimizers import Adam, RMSprop
from tensorflow.keras.preprocessing.image import ImageDataGenerator
from tensorflow.keras.callbacks import ModelCheckpoint
import numpy as np
import tensorflow as tf

from tensorflow.python.keras import models
from tensorflow.python.keras import layers

from tensorflow.keras import optimizers

from os import listdir

# data_list = listdir('/content/covid-19/four_classes/train')
data_list = listdir('coronet_org_data/four_classes/train')

# Delete some classes that may interfere

print(len(data_list))

start_time = time.time()

DATASET_PATH = 'coronet_org_data/four_classes/train'
test_dir = 'coronet_org_data/four_classes/test'
IMAGE_SIZE = (150, 150)
NUM_CLASSES = len(data_list)
BATCH_SIZE = 10  # try reducing batch size or freeze more layers if your GPU runs out of memory
NUM_EPOCHS = 80
LEARNING_RATE = 0.0001

# Train datagen here is a preprocessor
train_datagen = ImageDataGenerator(rescale=1. / 255,
                                   rotation_range=50,
                                   featurewise_center=True,
                                   featurewise_std_normalization=True,
                                   width_shift_range=0.2,
                                   height_shift_range=0.2,
                                   shear_range=0.25,
                                   zoom_range=0.1,
                                   zca_whitening=True,
                                   channel_shift_range=20,
                                   horizontal_flip=True,
                                   vertical_flip=True,
                                   validation_split=0.2,
                                   fill_mode='constant')

# For multiclass use categorical n for binary use binary
train_batches = train_datagen.flow_from_directory(DATASET_PATH,
                                                  target_size=IMAGE_SIZE,
                                                  shuffle=True,
                                                  batch_size=BATCH_SIZE,
                                                  subset="training",
                                                  seed=42,
                                                  class_mode="categorical"
                                                  # For multiclass use categorical n for binary use binary
                                                  )

valid_batches = train_datagen.flow_from_directory(DATASET_PATH,
                                                  target_size=IMAGE_SIZE,
                                                  shuffle=True,
                                                  batch_size=BATCH_SIZE,
                                                  subset="validation",
                                                  seed=42,
                                                  class_mode="categorical"
                                                  # For multiclass use categorical n for binary use binary

                                                  )

from tensorflow.keras.applications import Xception

conv_base = Xception(weights='imagenet',
                     include_top=False,
                     input_shape=(150, 150, 3))

conv_base.trainable = True

model = models.Sequential()
model.add(conv_base)

model.add(layers.Flatten())
model.add(layers.Dropout(0.5))
model.add(layers.Dense(256, activation='relu'))
model.add(layers.Dense(4, activation='softmax'))

model.compile(loss='categorical_crossentropy',  # for multiclass use categorical_crossentropy
              optimizer=optimizers.Adam(lr=LEARNING_RATE),
              metrics=['acc'])

print("Batch len")
print(len(train_batches))
print(len(valid_batches))

STEP_SIZE_TRAIN=train_batches.n//train_batches.batch_size
STEP_SIZE_VALID=valid_batches.n//valid_batches.batch_size

result=model.fit_generator(train_batches,
                        steps_per_epoch =STEP_SIZE_TRAIN,
                        validation_data = valid_batches,
                        validation_steps = STEP_SIZE_VALID,
                        epochs= NUM_EPOCHS,
                       )

print("Total time: "+str(time.time()-start_time))

import matplotlib.pyplot as plt


def plot_acc_loss(result, epochs):
    acc = result.history['acc']
    loss = result.history['loss']
    val_acc = result.history['val_acc']
    val_loss = result.history['val_loss']
    plt.figure(figsize=(15, 5))
    plt.subplot(121)
    plt.plot(range(1, epochs), acc[1:], label='Train_acc')
    plt.plot(range(1, epochs), val_acc[1:], label='Val_acc')
    plt.title('Accuracy over ' + str(epochs) + ' Epochs', size=15)
    plt.legend()
    plt.grid(True)
    plt.subplot(122)
    plt.plot(range(1, epochs), loss[1:], label='Train_loss')
    plt.plot(range(1, epochs), val_loss[1:], label='Val_loss')
    plt.title('Loss over ' + str(epochs) + ' Epochs', size=15)
    plt.legend()
    plt.grid(True)
    plt.show()


plot_acc_loss(result, 80)

# %%

# Save the trained model and copy to drive

model.save('4-class-Covid19-Mod-Xception.h5')
# !cp /content/"4-class-Covid19-Mod-Xception.h5" /content/drive/"My Drive"/"Colab Notebooks"


# %% md

# ** Evaluate
# using
# evaluate
# Generator **

# %%

# Create evaluate data generator from test set
# Dont forget shuffle false

test_datagen = ImageDataGenerator(rescale=1. / 255)
# test_dir = '/content/COVID-19 Radiography Database'
eval_generator = test_datagen.flow_from_directory(test_dir, target_size=IMAGE_SIZE, batch_size=1,
                                                  shuffle=False, seed=42, class_mode="categorical")
eval_generator.reset()

# %%

# Evalute the trained model on evaluate generator
eval_generator.reset()
x = model.evaluate_generator(eval_generator,
                             steps=np.ceil(len(eval_generator)),
                             use_multiprocessing=False,
                             verbose=1,
                             workers=1,
                             )

print('Test loss:', x[0])
print('Test accuracy:', x[1])

# Poor test accuracy due to the small dataset size

# %% md

# ** Create
# DataGen
# on
# single
# folder /
#
#
# class and predict ! **

# %%


# IMAGE_SIZE = (150, 150)
# test_datagen = ImageDataGenerator(rescale=1. / 255)
# test_dir = 'coronet_org_data/COVID-19 Radiography Database'
# pred_generator = test_datagen.flow_from_directory(
#     test_dir, target_size=IMAGE_SIZE,
#     batch_size=1,
#     shuffle=False,
#
#     seed=42,
#
#     class_mode="categorical")
# pred_generator.reset()
#
# count = [0, 0, 0, 0]
#
# files = pred_generator.filenames
#
# for i in range(len(files)):
#     x, y = pred_generator.next()
#     img = x
#     predict = model.predict(img)
#
#     p = np.argmax(predict, axis=-1)
#     print(str(p[0]) + " " + files[pred_generator.batch_index - 1])
#     # print(predict)
#     # p=model.predict_classes(img)
#     count[p[0]] += 1
#
# # print(str(p[0])+" "+files[i])
# print(count)
#
# # %% md
#
# ### **`Predict Results using predict generator and evaluate the accuracy and Confusion matrix `**
#
# # %%
#
# from sklearn.metrics import confusion_matrix
# from sklearn.metrics import plot_confusion_matrix
# from sklearn.metrics import classification_report
#
# filenames = eval_generator.filenames
# nb_samples = len(filenames)
# eval_generator.reset()
# predict = model.predict_generator(eval_generator, steps=np.ceil(len(eval_generator)))
# pp = predict
# predict = np.argmax(predict, axis=-1)
# classes = eval_generator.classes[eval_generator.index_array]
# acc = sum(predict == classes) / len(predict)
# names = ["covid", "normal", "pneumonia_bac", "pneumonia_vir"]
# # print(confusion_matrix(classes,predict))
#
# font = {
#     'family': 'Times New Roman',
#     'size': 12
# }
# plt.rc('font', **font)
# cm = confusion_matrix(classes, predict)
# print(cm)
# print(classification_report(classes, predict))
# plt.imshow(cm, cmap=plt.cm.Blues)
# plt.xlabel('Predicted labels \nAccuracy: {:0.2f}'.format(acc * 100))
# plt.ylabel("True labels")
# plt.xticks(classes, [])
# plt.yticks(classes, [])
# plt.title('Confusion matrix ')
# plt.colorbar()
# plt.show()

# %% md

# ** Test
# Single
# image **

# %%

# import cv2
# from skimage import transform
#
# img_r = cv2.imread('/content/test/x.jpg')
#
# img1 = np.array(img_r).astype('float32') / 255
# img2 = transform.resize(img1, (150, 150, 3))
#
# img = np.expand_dims(img2, axis=0)
#
# r = model.predict(img)
#
# names = dict((v, k) for k, v in labels.items())
# index = np.argmax(r)
# name = names.get(index, "Unknown")
#
# p = round(r.max() * 100, 3)  # to find maximum score
#
# scores = r
# print(scores)
#
# font = {
#     'family': 'Times New Roman',
#     'size': 9,
#
# }
# plt.rc('font', **font)
#
# # plt.title(name +" ("+ str(p)+")")
# plt.title(names[0] + " " + str(round(scores[0][0] * 100, 1)) + "%" + "\n" + names[1] + " " + str(
#     round(scores[0][1] * 100, 1)) + "%" + "\n" + names[2] + " " + str(round(scores[0][2] * 100, 1)) + "%" + "\n" +
#           names[3] + " " + str(round(scores[0][3] * 100, 1)) + "%")
#
# plt.imshow(img2)
#
# # %% md
#
# # ** Test
# # Whole
# # Folder **
#
# # %%
#
# import cv2
# from skimage import transform
#
# count = [0, 0, 0, 0]
# folder_name = "/content/drive/My Drive/Datasets/covid-19/covidnew/covid"
# files = os.listdir(folder_name)
# for i in range(len(files)):
#     img_r = cv2.imread(folder_name + "/" + files[i])
#
#     img = np.array(img_r).astype('float32') / 255
#
#     img = transform.resize(img, (150, 150, 3))
#     img = np.expand_dims(img, axis=0)
#
#     predict = model.predict(img)
#     p = np.argmax(predict, axis=-1)
#     # p=model.predict_classes(img)
#     count[p[0]] += 1
#     print(str(p[0]) + " " + files[i])
#
#     print()
#
# print(count)