import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers
import numpy as np
import cv2


class ImageClassifier:
    def __init__(self):
        self.model = None

    def preprocess_image(self, image):
        # Resize the image to (32, 32)
        resized_image = cv2.resize(image, (32, 32))

        # # Convert the image to grayscale
        # gray_image = cv2.cvtColor(resized_image, cv2.COLOR_BGR2GRAY)

        # # # Normalize the pixel values between 0 and 1
        # normalized_image = gray_image.astype("float32") / 255.0

        # # # Transpose the dimensions to match the model's input shape
        # transposed_image = np.transpose(normalized_image, (1, 2, 0))

        # # # Expand dimensions to match model input shape (add batch dimension)
        # img_array = np.expand_dims(transposed_image, axis=0)
        return resized_image

    def load_dataset(self):
        # Set up the dataset
        (x_train, y_train), (x_test, y_test) = keras.datasets.cifar10.load_data()

        # Normalize pixel values between 0 and 1
        x_train = x_train.astype("float32") / 255.0
        x_test = x_test.astype("float32") / 255.0

        return (x_train, y_train), (x_test, y_test)

    # def build_model(self, x_train):
    #     # Define the model architecture
    #     model = keras.Sequential([
    #         # keras.Input(shape=x_train.shape[1]),
    #         layers.Conv2D(32, kernel_size=(3, 3), activation="relu", padding='same'),
    #         layers.MaxPooling2D(pool_size=(2, 2)),
    #         layers.Conv2D(64, kernel_size=(3, 3), activation="relu", padding='same'),
    #         layers.MaxPooling2D(pool_size=(2, 2)),
    #         layers.Flatten(),
    #         layers.Dropout(0.5),
    #         layers.Dense(10, activation="softmax")
    #     ])

    #     # Compile the model
    #     model.compile(loss="sparse_categorical_crossentropy", optimizer="adam", metrics=["accuracy"])

    #     self.model = model
    
    def build_model(self, x_train):
      # Define the model architecture
      model = keras.Sequential([
          layers.Conv2D(32, kernel_size=(3, 3), activation="relu", padding='same'),
          layers.BatchNormalization(),
          layers.MaxPooling2D(pool_size=(2, 2)),
          layers.Dropout(0.25),

          layers.Conv2D(64, kernel_size=(3, 3), activation="relu", padding='same'),
          layers.BatchNormalization(),
          layers.MaxPooling2D(pool_size=(2, 2)),
          layers.Dropout(0.25),

          layers.Conv2D(128, kernel_size=(3, 3), activation="relu", padding='same'),
          layers.BatchNormalization(),
          layers.MaxPooling2D(pool_size=(2, 2)),
          layers.Dropout(0.25),

          layers.Flatten(),
          layers.Dense(256, activation="relu"),
          layers.BatchNormalization(),
          layers.Dropout(0.5),

          layers.Dense(10, activation="softmax")
      ])

      # Compile the model
      optimizer = keras.optimizers.RMSprop(learning_rate=0.001)
      model.compile(loss="sparse_categorical_crossentropy", optimizer=optimizer, metrics=["accuracy"])

      self.model = model

    def train_model(self, x_train, y_train, batch_size, epochs, validation_split):
        # Train the model
        self.model.fit(x_train, y_train, batch_size=batch_size, epochs=epochs, verbose=1, validation_split=validation_split)

    def evaluate_model(self, x_test, y_test):
        # Evaluate the model on the test set
        score = self.model.evaluate(x_test, y_test, verbose=0)
        print("Test loss:", score[0])
        print("Test accuracy:", score[1])

    def save_model(self, filepath):
        # Save the trained model
        self.model.save(filepath)

    def load_model(self, filepath):
        # Load the trained model
        self.model = keras.models.load_model(filepath)

    def classify_image(self, image, top_k=3):
        # Preprocess the image
        preprocessed_image = self.preprocess_image(image)

        # Perform inference
        predicted_probs = self.model.predict(np.array([preprocessed_image]))
        top_classes = np.argsort(predicted_probs[0])[-top_k:][::-1]
        top_probs = predicted_probs[0][top_classes]

        return top_classes, top_probs