#Tensorflow training two

import warnings

warnings.filterwarnings('ignore')

import cv2
import numpy as np
import os
import matplotlib.pyplot as plt
from sklearn.model_selection import train_test_split
from keras.preprocessing.image import ImageDataGenerator
from keras.utils.np_utils import to_categorical
from keras.models import Sequential
from keras.layers import Conv2D, Dense, MaxPooling2D, Activation, Dropout, Flatten
from tensorflow.keras.optimizers import Adam
import mplcursors
from tensorflow.keras.applications import MobileNetV2
from tensorflow.keras.layers import Input
from tensorflow.keras.applications.mobilenet_v2 import preprocess_input
########################################

path = r'C:\Users\adamt\Desktop\archive'

# Saves the Photos that the data gen generates to view before training
# This could cause issues so I left the original code in its original place

brightness = [0.1, 0.5]
width = [0.01, 0.05]
height = [0.01, 0.05]



images = []
classNo = []
testRatio = 0.2
valRatio = 0.2
imgDimension = (167, 167, 3)  # 1 is Black and white, 3 is RGB
img_counter = 0
#########################################

myList = os.listdir(path)

numOfClasses = len(myList)

print(numOfClasses)

print("Importing Classes..........")
for x in range(0, numOfClasses):
    myPicList = os.listdir(path + "/" + str(x))
    # myData/0/img.jpg # Not sure what this does
    for y in myPicList:
        curImg = cv2.imread(path + "/" + str(x) + "/" + y)
        curImg = cv2.resize(curImg, (imgDimension[0], imgDimension[1]))
        images.append(curImg)
        classNo.append(x)
    print(x, "X")

images = np.array(images)
classNo = np.array(classNo)

#########Spliting The Data###########

x_train, x_test, y_train, y_test = train_test_split(images, classNo, test_size=testRatio)
x_train, x_validation, y_train, y_validation = train_test_split(x_train, y_train, test_size=valRatio)

print(x_train.shape)
#print(x_validation,"x_validation")
#print(y_validation,"y_validation")
#print(testRatio,"test Ratio")
numOfSample = []

for x in range(0, numOfClasses):
    numOfSample.append(len(np.where(y_train == x)[0]))

plt.figure(figsize=(10, 5))
plt.bar(range(0, numOfClasses), numOfSample)
# print(numOfSample,"number of samples")
plt.title("Bar Plot of Classes & Images")
plt.xlabel("No Of Classes")
plt.ylabel("No of Images")
mplcursors.cursor(hover=True)  # Hover over bar charts tells you how many photos are being used(y)
plt.show()


def preprocessing(img):
    # img=np.astype("uint8")
    img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    #cv2.imshow("v",img)
    img = cv2.equalizeHist(img)
    img = img / 255
    return img


x_train = np.array(list(map(preprocessing, x_train)))
x_test = np.array(list(map(preprocessing, x_test)))
x_validation = np.array(list(map(preprocessing, x_validation)))

x_train = x_train.reshape(x_train.shape[0], x_train.shape[1], x_train.shape[2], 1)
x_test = x_test.reshape(x_test.shape[0], x_test.shape[1], x_test.shape[2], 1)
x_validation = x_validation.reshape(x_validation.shape[0], x_validation.shape[1], x_validation.shape[2], 1)

# brightness = [0.1, 0.5]
# width = [0.01, 0.05]
# height = [0.01, 0.05]

dataGen = ImageDataGenerator(
    featurewise_center=True,
    featurewise_std_normalization=True,
    rotation_range=20,
    width_shift_range=0.02,
    height_shift_range=0.02,
    horizontal_flip=True,
    brightness_range=brightness,
    validation_split=0.90)

# Save Augmented Photos
# i = 0
# for batch in dataGen.flow_from_directory(directory=path,
#                                         batch_size=2,
#                                         save_to_dir=r'C:\Users\adamt\Desktop\dataset\New_one_pence_two_Pence\2',
#                                         save_prefix='Coins_agumented_from_trainer_2',
#                                         target_size=(224, 224),
#                                         save_format='png'):
#
#    i += 1
#    print(i, "i")
#    if i > 31:
#        break

y_train = to_categorical(y_train, numOfClasses)
y_test = to_categorical(y_test, numOfClasses)
y_validation = to_categorical(y_validation, numOfClasses)

baseModel = MobileNetV2(weights="imagenet", include_top=False,
    input_tensor=Input(shape=(128,129,3)))

def myModel():
    sizeOfFilter1 = (3, 3)
    sizeOfFilter2 = (3, 3)
    sizeOfPool = (2, 2)

    model = Sequential()
    model.add((Conv2D(32, sizeOfFilter1, input_shape=(imgDimension[0], imgDimension[1], 1), activation='relu')))
    model.add((Conv2D(32, sizeOfFilter1, activation='relu')))
    model.add(MaxPooling2D(pool_size=sizeOfPool))

    model.add((Conv2D(64, sizeOfFilter2, activation='relu')))
    model.add((Conv2D(64, sizeOfFilter2, activation='relu')))
    model.add(MaxPooling2D(pool_size=sizeOfPool))
    model.add(Dropout(0.8))  # Changed from 0.5 to 0.8

    model.add(Flatten())
    model.add(Dense(64, activation='relu'))
    model.add(Dropout(0.5))
    model.add(Dense(numOfClasses, activation='softmax'))
    model.compile(Adam(lr=0.001), loss='categorical_crossentropy', metrics=['accuracy'])
    return model


model = myModel()
print(model.summary())

history = model.fit_generator(dataGen.flow(x_train, y_train, batch_size=32),
                              steps_per_epoch=100,
                              # https://chowdera.com/2022/01/202201090452117647.html amount of images divided by
                              # batch size, only need to do for 1 folder
                              epochs=5, validation_data=(x_validation, y_validation), shuffle=250)
model.save("Eye_Detector_100_Train.h5")