import os
import numpy as np
import tensorflow as tf
from tensorflow import keras
import cv2
import random

#YOU CAN DOWNLOAD MODEL FROM https://uam-my.sharepoint.com/:f:/g/personal/pavbia_st_amu_edu_pl/EmBHjnETuk5LiCZS6xk7AnIBNsnffR3Sygf8EX2bhR1w4A

class VacuumRecognizer:
    model = keras.models.load_model('AI_brain\model.h5')

    def recognize(self, image_path) -> str:
        class_names = ['Banana', 'Cat', 'Earings', 'Plant']

        img = cv2.imread(image_path, flags=cv2.IMREAD_GRAYSCALE)
        cv2.waitKey(0)
        img = (np.expand_dims(img, 0))

        predictions = self.model.predict(img)[0].tolist()

        # More information about neuron decision
        # print(img.shape)
        # cv2.imshow("test_show", img)
        # print(class_names)
        # print(predictions)
        # print(max(predictions))
        # print(predictions.index(max(predictions)))

        return class_names[predictions.index(max(predictions))]

    def get_random_dir(self, type) -> str:
        if type == 'Plant':
            plant_image_paths = 'C:\\Users\\Pavel\\Desktop\\AI\\Machine_learning_2023\\AI_brain\\Image_datasetJPGnewBnW\\check\\Plant'
            plant_dirs = os.listdir(plant_image_paths)
            full_path = plant_image_paths + '\\' + plant_dirs[random.randint(0, len(plant_dirs)-1)]
            print(full_path)
            return full_path
        elif type == 'Earings':
            earnings_image_paths = 'C:\\Users\\Pavel\\Desktop\\AI\\Machine_learning_2023\\AI_brain\\Image_datasetJPGnewBnW\\check\\Earings'
            earning_dirs = os.listdir(earnings_image_paths)
            full_path = earnings_image_paths + '\\' + earning_dirs[random.randint(0, len(earning_dirs)-1)]
            print(full_path)
            return full_path
        elif type == 'Banana':
            banana_image_paths = 'C:\\Users\\Pavel\\Desktop\\AI\\Machine_learning_2023\\AI_brain\\Image_datasetJPGnewBnW\\check\\Banana'
            banana_dirs = os.listdir(banana_image_paths)
            full_path = banana_image_paths + '\\' + banana_dirs[random.randint(0, len(banana_dirs)-1)]
            print(full_path)
            return full_path
        elif type == 'Cat':
            cat_image_paths = 'C:\\Users\\Pavel\\Desktop\\AI\\Machine_learning_2023\\AI_brain\\Image_datasetJPGnewBnW\\check\\Cat'
            cat_dir = os.listdir(cat_image_paths)


#For testing the neuron model
'''image_paths = []
image_paths.append('C:\\Users\\Pavel\\Desktop\\AI\\Machine_learning_2023\\AI_brain\\Image_datasetJPGnewBnW\\check\\Banana')
image_paths.append('C:\\Users\\Pavel\\Desktop\\AI\\Machine_learning_2023\\AI_brain\\Image_datasetJPGnewBnW\\check\\Cat')
image_paths.append('C:\\Users\\Pavel\\Desktop\\AI\\Machine_learning_2023\\AI_brain\\Image_datasetJPGnewBnW\\check\\Earings')
image_paths.append('C:\\Users\\Pavel\\Desktop\\AI\\Machine_learning_2023\\AI_brain\\Image_datasetJPGnewBnW\\check\\Plant')
uio = VacuumRecognizer()

for image_path in image_paths:
    dirs = os.listdir(image_path)
    for i in range(3):
        print(uio.recognize(image_path + '\\' + dirs[random.randint(0, len(dirs)-1)]))'''