import os
from os.path import basename, join, exists
import numpy as np
import math
np.random.seed(777)
import tensorflow_addons as tfa
from tensorflow.keras import models
from tensorflow.keras import layers 
from tensorflow.keras import models
from tensorflow.keras import optimizers
from tensorflow.keras import callbacks
from tensorflow.keras import callbacks
from tensorflow.keras.applications.xception import Xception
from tensorflow.keras.preprocessing.image import ImageDataGenerator
from numpy import array
from numpy import argmax
from  numpy import mean 
from numpy import std
from sklearn.metrics import accuracy_score
import matplotlib.pyplot as plt
import warnings
warnings.filterwarnings('always')
warnings.filterwarnings('ignore')


os.chdir(r"COVID_Xray/")
train_dir="aug/"
test_dir="test/"

total=0
print('---Training set details----')
for sub_folder in os.listdir(train_dir):
  no_of_images=len(os.listdir(train_dir + sub_folder))
  total+=no_of_images
  print(str(no_of_images) + " " + sub_folder + " images")

print("Total no. of Chest Xray training images=",total)
total=0
print('---Test set details----')
for sub_folder in os.listdir(test_dir):
  no_of_images=len(os.listdir(test_dir + sub_folder))
  total+=no_of_images
  print(str(no_of_images) + " " + sub_folder + " images")

print("Total no. of Chest Xray test images=",total)

extracted_features_dir="COVID_Xray/extracted_features/"

img_height =512
img_width = 512
batch_size =32
input_shape = (img_width, img_height, 3)

print("-----------------Image Augmentation for Xception--------------")

random_seed = np.random.seed(1142)
train_datagen = ImageDataGenerator(
    rescale=1./255,
    featurewise_center=True,
    featurewise_std_normalization=True,
    rotation_range=15,
    width_shift_range=0.2,
    height_shift_range=0.2,
    validation_split= 0.2,
    zoom_range=0.1,
    shear_range=0.2)

train_generator_xcep = train_datagen.flow_from_directory(
    train_dir,
    target_size=(img_height, img_width),
    batch_size=batch_size,
    seed = random_seed,
    shuffle=False,
    subset = 'training',
    class_mode='categorical')

val_generator_xcep = train_datagen.flow_from_directory(
    train_dir,
    target_size=(img_height, img_width),
    batch_size=batch_size,
    seed = random_seed,
    shuffle=False,
    subset = 'validation',
    class_mode='categorical')

test_datagen=ImageDataGenerator(rescale=1./255)
test_generator_xcep=test_datagen.flow_from_directory(test_dir,
                                                      target_size=(img_height, img_width),
                                                          batch_size=batch_size, 
                                                          seed=random_seed,
                                                          shuffle=False,
                                                          class_mode='categorical') 

nb_train_samples = len(train_generator_xcep.filenames)
nb_validation_samples = len(val_generator_xcep.filenames)
predict_size_train = int(math.ceil(nb_train_samples / batch_size))
predict_size_validation = int(math.ceil(nb_validation_samples / batch_size))
nb_test_samples = len(test_generator_xcep.filenames)
predict_size_test = int(math.ceil(nb_test_samples / batch_size))

model_name="Xception"
model = Xception(include_top=False, weights="imagenet",pooling='avg',input_shape=input_shape)
image_input =model.input
x1 = layers.GlobalAveragePooling2D()(model.get_layer("block4_sepconv1").output) #layer_27
x2 = layers.GlobalAveragePooling2D()(model.get_layer("block5_sepconv1").output) #layer 37
x3 = layers.GlobalAveragePooling2D()(model.get_layer("block14_sepconv1").output)  #layer_126
out= layers.Concatenate()([x1,x2,x3])
out=layers.Dense(512,activation='relu')(out)
out=layers.Dropout(0.5)(out)
out=layers.Dense(3,activation='softmax',name= 'output')(out)
custom_xcep_model = models.Model(image_input , out)
custom_xcep_model.summary()

for layer in custom_xcep_model.layers[:115]:
    layer.trainable = False
custom_xcep_model.summary()

nEpochs=100
base_lr=1e-3

opt = optimizers.Adam(lr=base_lr, beta_1=0.6, beta_2=0.8,amsgrad=True)
custom_xcep_model.compile(optimizer = opt, loss='categorical_crossentropy', metrics=['accuracy'])
checkpoint1 = callbacks.ModelCheckpoint('saved models/Xception/xception_weights.h5', monitor='val_accuracy', verbose=1, save_best_only=True, mode='max')
callbacks_list=[checkpoint1]
# Training the modified Xception network for refining the deep feature embedding
history =custom_xcep_model.fit(train_generator_xcep,
                    epochs=nEpochs,
                    validation_data=val_generator_xcep,
                    callbacks=callbacks_list)
bottleneck= tf.keras.Model(inputs=custom_xcep_model.input, outputs=custom_xcep_model.layers[130].output)
#Saving features of the training images
features_train =bottleneck.predict_generator(train_generator_xcep, predict_size_train)
np.save(extracted_features_dir+model_name+'_train_features.npy', features_train)

# Saving features of the validation images
features_validation =bottleneck.predict_generator(val_generator_xcep, predict_size_validation)
np.save(extracted_features_dir+model_name+'_val_features.npy', features_validation)

# Saving features of the test images
features_test = bottleneck.predict_generator(test_generator_xcep, predict_size_test)
np.save(extracted_features_dir+model_name+'_test_features.npy', features_test)

tensorflow.keras.layers.Concatenate tensorflow.keras.callbacks.ModelCheckpoint tensorflow.keras.layers.GlobalAveragePooling2D tensorflow.keras.preprocessing.image.ImageDataGenerator tensorflow.keras.applications.xception.Xception numpy.random.seed tensorflow.keras.models.Model tensorflow.keras.layers.Dense numpy.save tensorflow.keras.optimizers.Adam tensorflow.keras.layers.Dropout
[(5, 'numpy.random.seed', 'np.random.seed', (['(777)'], {}), True, 'import numpy as np\n'), (22, 'warnings.filterwarnings', 'warnings.filterwarnings', (['"""always"""'], {}), False, 'import warnings\n'), (23, 'warnings.filterwarnings', 'warnings.filterwarnings', (['"""ignore"""'], {}), False, 'import warnings\n'), (26, 'os.chdir', 'os.chdir', (['"""COVID_Xray/"""'], {}), False, 'import os\n'), (32, 'os.listdir', 'os.listdir', (['train_dir'], {}), False, 'import os\n'), (40, 'os.listdir', 'os.listdir', (['test_dir'], {}), False, 'import os\n'), (56, 'numpy.random.seed', 'np.random.seed', (['(1142)'], {}), True, 'import numpy as np\n'), (57, 'tensorflow.keras.preprocessing.image.ImageDataGenerator', 'ImageDataGenerator', ([], {'rescale': '(1.0 / 255)', 'featurewise_center': '(True)', 'featurewise_std_normalization': '(True)', 'rotation_range': '(15)', 'width_shift_range': '(0.2)', 'height_shift_range': '(0.2)', 'validation_split': '(0.2)', 'zoom_range': '(0.1)', 'shear_range': '(0.2)'}), False, 'from tensorflow.keras.preprocessing.image import ImageDataGenerator\n'), (86, 'tensorflow.keras.preprocessing.image.ImageDataGenerator', 'ImageDataGenerator', ([], {'rescale': '(1.0 / 255)'}), False, 'from tensorflow.keras.preprocessing.image import ImageDataGenerator\n'), (102, 'tensorflow.keras.applications.xception.Xception', 'Xception', ([], {'include_top': '(False)', 'weights': '"""imagenet"""', 'pooling': '"""avg"""', 'input_shape': 'input_shape'}), False, 'from tensorflow.keras.applications.xception import Xception\n'), (111, 'tensorflow.keras.models.Model', 'models.Model', (['image_input', 'out'], {}), False, 'from tensorflow.keras import models\n'), (121, 'tensorflow.keras.optimizers.Adam', 'optimizers.Adam', ([], {'lr': 'base_lr', 'beta_1': '(0.6)', 'beta_2': '(0.8)', 'amsgrad': '(True)'}), False, 'from tensorflow.keras import optimizers\n'), (123, 'tensorflow.keras.callbacks.ModelCheckpoint', 'callbacks.ModelCheckpoint', (['"""saved models/Xception/xception_weights.h5"""'], {'monitor': '"""val_accuracy"""', 'verbose': '(1)', 'save_best_only': '(True)', 'mode': '"""max"""'}), False, 'from tensorflow.keras import callbacks\n'), (133, 'numpy.save', 'np.save', (["(extracted_features_dir + model_name + '_train_features.npy')", 'features_train'], {}), True, 'import numpy as np\n'), (137, 'numpy.save', 'np.save', (["(extracted_features_dir + model_name + '_val_features.npy')", 'features_validation'], {}), True, 'import numpy as np\n'), (141, 'numpy.save', 'np.save', (["(extracted_features_dir + model_name + '_test_features.npy')", 'features_test'], {}), True, 'import numpy as np\n'), (96, 'math.ceil', 'math.ceil', (['(nb_train_samples / batch_size)'], {}), False, 'import math\n'), (97, 'math.ceil', 'math.ceil', (['(nb_validation_samples / batch_size)'], {}), False, 'import math\n'), (99, 'math.ceil', 'math.ceil', (['(nb_test_samples / batch_size)'], {}), False, 'import math\n'), (104, 'tensorflow.keras.layers.GlobalAveragePooling2D', 'layers.GlobalAveragePooling2D', ([], {}), False, 'from tensorflow.keras import layers\n'), (105, 'tensorflow.keras.layers.GlobalAveragePooling2D', 'layers.GlobalAveragePooling2D', ([], {}), False, 'from tensorflow.keras import layers\n'), (106, 'tensorflow.keras.layers.GlobalAveragePooling2D', 'layers.GlobalAveragePooling2D', ([], {}), False, 'from tensorflow.keras import layers\n'), (107, 'tensorflow.keras.layers.Concatenate', 'layers.Concatenate', ([], {}), False, 'from tensorflow.keras import layers\n'), (108, 'tensorflow.keras.layers.Dense', 'layers.Dense', (['(512)'], {'activation': '"""relu"""'}), False, 'from tensorflow.keras import layers\n'), (109, 'tensorflow.keras.layers.Dropout', 'layers.Dropout', (['(0.5)'], {}), False, 'from tensorflow.keras import layers\n'), (110, 'tensorflow.keras.layers.Dense', 'layers.Dense', (['(3)'], {'activation': '"""softmax"""', 'name': '"""output"""'}), False, 'from tensorflow.keras import layers\n'), (33, 'os.listdir', 'os.listdir', (['(train_dir + sub_folder)'], {}), False, 'import os\n'), (41, 'os.listdir', 'os.listdir', (['(test_dir + sub_folder)'], {}), False, 'import os\n')]