import os

import matplotlib.pyplot as plt
import numpy as np
import tensorflow as tf

from tensorflow import keras
from tensorflow.keras import layers
from tensorflow.keras.models import Sequential

import pathlib

from tensorflow.python.client import device_lib
print(device_lib.list_local_devices())

data_dir = "C:/Users/jilek/Downloads/AAT+"
data_dir = pathlib.Path(data_dir).with_suffix('')

data_dir_test = "C:/Users/jilek/Downloads/AAT+_TEST"
data_dir_test = pathlib.Path(data_dir_test).with_suffix('')

image_count = len(list(data_dir.glob('*/*.jpg')))
print(image_count)

batch_size = 1
img_height = 1024
img_width = 1024

train_ds = tf.keras.utils.image_dataset_from_directory(
    data_dir,
    validation_split=0.0,
    #subset="training",
    seed=123,
    labels='inferred',
    label_mode='categorical',
    class_names=["C100", "C095", "C090", "C085", "C080", "C070", "C060", "C040", "C020"],
    color_mode="grayscale", #grayscale
    shuffle=True,
    image_size=(img_height, img_width),
    batch_size=batch_size)

val_ds = tf.keras.utils.image_dataset_from_directory(
    data_dir_test,
    validation_split=0.0,
    #subset="validation",
    seed=123,
    labels='inferred',
    label_mode='categorical',
    class_names=["C100", "C095", "C090", "C085", "C080", "C070", "C060", "C040", "C020"],
    color_mode="grayscale",
    image_size=(img_height, img_width),
    batch_size=batch_size)

class_names = train_ds.class_names
print(class_names)

for image_batch, labels_batch in train_ds:
    print(image_batch.shape)
    print(labels_batch.shape)
    break

AUTOTUNE = tf.data.AUTOTUNE

data_augmentation = keras.Sequential(
    [
        layers.RandomFlip("horizontal_and_vertical",
                          input_shape=(img_height,
                                       img_width,
                                       1)), #rgb
        #layers.RandomRotation(0.5),
        #layers.RandomZoom(0.5),
    ]
)

train_ds = train_ds.shuffle(buffer_size=900).prefetch(buffer_size=AUTOTUNE) #.cache()
val_ds = val_ds.prefetch(buffer_size=AUTOTUNE) #.cache()


num_classes = len(class_names)
print(str(num_classes))

model = Sequential([
  layers.Rescaling(1.0/255, input_shape=(img_height, img_width, 1)), #rgb
  #layers.Dropout(0.0),
  #layers.MaxPooling2D(pool_size=(8, 8)),
  layers.Conv2D(4, (4, 4), strides=(2, 2), padding='valid', dilation_rate=(1, 1), groups=1, input_shape=(1024, 1024, 1), activation='relu'),
  layers.Conv2D(8, (4, 4), strides=(2, 2), padding='valid', dilation_rate=(1, 1), groups=1, input_shape=(512, 512, 4), activation='relu'),
  layers.Conv2D(16, (4, 4), strides=(4, 4), padding='valid', dilation_rate=(1, 1), groups=1, input_shape=(256, 256, 8), activation='relu'),
  layers.Conv2D(32, (4, 4), strides=(4, 4), padding='valid', dilation_rate=(1, 1), groups=1, input_shape=(64, 64, 16), activation='relu'),
  layers.Conv2D(64, (4, 4), strides=(4, 4), padding='valid', dilation_rate=(1, 1), groups=1, input_shape=(16, 16, 32), activation='relu'),
  #layers.Conv2D(128, (4, 4), strides=(1, 1), padding='valid', dilation_rate=(1, 1), groups=1, input_shape=(8, 8, 64), activation='relu'),
  #layers.Dropout(0.1),
  layers.Flatten(),
  layers.Dense(32, activation='relu'),
  layers.Dense(num_classes, activation='softmax')
])
model.compile(optimizer=tf.keras.optimizers.Adam(learning_rate=1e-5),
              loss=tf.keras.losses.CategoricalCrossentropy(),
              metrics=['accuracy'])

model.summary()
model.save("./model/AAT+")

epochs = 130
history = model.fit(
  train_ds,
  validation_data=val_ds,
  epochs=epochs
)

acc = history.history['accuracy']
val_acc = history.history['val_accuracy']

loss = history.history['loss']
val_loss = history.history['val_loss']

epochs_range = range(epochs)

plt.figure(figsize=(8, 8))
plt.subplot(1, 2, 1)
plt.plot(epochs_range, acc, label='Training Accuracy')
plt.plot(epochs_range, val_acc, label='Validation Accuracy')
plt.legend(loc='lower right')
plt.title('Training and Validation Accuracy')

plt.subplot(1, 2, 2)
plt.plot(epochs_range, loss, label='Training Loss')
plt.plot(epochs_range, val_loss, label='Validation Loss')
plt.legend(loc='upper right')
plt.title('Training and Validation Loss')
plt.show()

test_dir = "C:/Users/jilek/Downloads/AAT_T/"
for file_name in os.listdir(test_dir):
    file_path = os.path.join(test_dir, file_name)
    img = tf.keras.utils.load_img(
        file_path, target_size=(img_height, img_width), color_mode="grayscale" #grayscale
    )
    img_array = tf.keras.utils.img_to_array(img)
    img_array = tf.expand_dims(img_array, 0)  # Create a batch

    predictions = model.predict(img_array)
    score = tf.nn.softmax(predictions[0])

    print(file_name)
    print(
        "This image most likely belongs to {} with a {:.2f} percent confidence."
        .format(class_names[np.argmax(score)], 100 * np.max(score))
    )