import os
import glob
import PIL
from PIL import Image
import tensorflow as tf
import pickle
from tensorflow import keras
from keras import layers
from keras.models import Sequential
import pathlib


class NeuralN:
    # @staticmethod
    def predict(self, image_path):
        data_dir = pathlib.Path('zdjecia')
        saved_model_path = pathlib.Path('trained_model.h5')
        class_names_path = pathlib.Path("class_names.pkl")
        image_count = sum(len(files) for _, _, files in os.walk(data_dir))

        print(image_count)

        # ORK_ARCHER = list(glob.glob('C:\\mobs_photos\\ORK_ARCHER'))
        # im = PIL.Image.open(ORK_ARCHER[0])
        # im.show()
        if os.path.exists(saved_model_path):
            model = tf.keras.models.load_model(saved_model_path)
            print("Saved model loaded")
            with open(class_names_path, 'rb') as f:
                class_names = pickle.load(f)
                print("Class names loaded.")

        else:
            train_ds = tf.keras.utils.image_dataset_from_directory(
                data_dir,
                validation_split=0.2,
                subset="training",
                seed=123,
                image_size=(180, 180),
                batch_size=32)

            val_ds = tf.keras.utils.image_dataset_from_directory(
                data_dir,
                validation_split=0.2,
                subset="validation",
                seed=123,
                image_size=(180, 180),
                batch_size=32)

            class_names = train_ds.class_names
            print(class_names)

            num_classes = len(class_names)

            model = Sequential([
                layers.Rescaling(1. / 255, input_shape=(180, 180, 3)),
                layers.Conv2D(16, 3, padding='same', activation='relu'),
                layers.MaxPooling2D(),
                layers.Conv2D(32, 3, padding='same', activation='relu'),
                layers.MaxPooling2D(),
                layers.Conv2D(64, 3, padding='same', activation='relu'),
                layers.MaxPooling2D(),
                layers.Flatten(),
                layers.Dense(128, activation='relu'),
                layers.Dense(num_classes)
            ])

            model.compile(optimizer='adam',
                          loss=tf.keras.losses.SparseCategoricalCrossentropy(
                              from_logits=True),
                          metrics=['accuracy'])
            model.summary()

            epochs = 10
            history = model.fit(
                train_ds,
                validation_data=val_ds,
                epochs=epochs
            )
            model.save("trained_model.h5")
            print("Model trained and saved.")
            with open(class_names_path, 'wb') as f:
                pickle.dump(train_ds.class_names, f)
                print("Class names saved.")
            # loaded_model = tf.keras.models.load_model("trained_model.h5")
        probability_model = tf.keras.Sequential([model,
                                                 tf.keras.layers.Softmax()])

        # image_path = image
        #image_path = pathlib.Path('')
        image = Image.open(image_path)

        # Preprocess the image

        image = image.resize((180, 180))  # Resize to match the input size of the model
        image_array = tf.keras.preprocessing.image.img_to_array(image)
        image_array = image_array / 255.0  # Normalize pixel values

        # Add an extra dimension to the image array
        image_array = tf.expand_dims(image, 0)
        # Make the prediction
        model = tf.keras.models.load_model("trained_model.h5")
        prediction = model.predict(image_array)
        # Convert the predictions to class labels
        predicted_label = class_names[prediction[0].argmax()]

        # Print the predicted label
        print(predicted_label)
        return predicted_label