import os
import json
import numpy as np
import tensorflow as tf
from tensorflow.keras.preprocessing.image import ImageDataGenerator
from sklearn.model_selection import train_test_split
from tensorflow.keras.applications import MobileNetV2
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Flatten, Dense

# Define image dimensions and channels
image_width, image_height, channels = 224, 224, 3

# Load class names from classes.json
with open("classes.json", "r") as f:
    class_names = json.load(f)

# Initialize the data generator with data augmentation
datagen = ImageDataGenerator(
    rescale=1.0 / 255.0,  # Normalize pixel values to the range [0, 1]
    rotation_range=20,   # Randomly rotate images
    width_shift_range=0.2,  # Randomly shift images horizontally
    height_shift_range=0.2,  # Randomly shift images vertically
    shear_range=0.2,  # Randomly apply shear transformations
    zoom_range=0.2,  # Randomly zoom in on images
    horizontal_flip=True,  # Randomly flip images horizontally
    fill_mode='nearest'  # Fill in missing pixels with the nearest value
)

# Load images and labels into arrays
images = []
labels = []

for class_num, class_name in class_names.items():
    class_path = os.path.join("dataset", class_num)
    for image_file in os.listdir(class_path):
        if image_file.endswith(".jpg"):  # Assuming your images are in JPG format
            image = tf.keras.preprocessing.image.load_img(
                os.path.join(class_path, image_file),
                target_size=(image_width, image_height)
            )
            image = tf.keras.preprocessing.image.img_to_array(image)
            images.append(image)
            labels.append(int(class_num))

# Convert lists to numpy arrays
images = np.array(images)
labels = np.array(labels)

# Split the dataset into training and testing sets
X_train, X_test, y_train, y_test = train_test_split(images, labels, test_size=0.2, random_state=42)

# Load MobileNetV2 as a feature extractor (include_top=False)
base_model = MobileNetV2(weights='imagenet', include_top=False, input_shape=(image_width, image_height, channels))

# Freeze the layers of MobileNetV2
for layer in base_model.layers:
    layer.trainable = False

# Create your custom classification head
model = Sequential([
    base_model,
    Flatten(),
    Dense(128, activation='relu'),
    Dense(len(class_names), activation='softmax')
])

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

# Train the model using data augmentation
history = model.fit(datagen.flow(X_train, y_train, batch_size=32), epochs=37, validation_data=(X_test, y_test))

# Evaluate the model on the test data
test_loss, test_acc = model.evaluate(X_test, y_test, verbose=2)
print("\nTest accuracy:", test_acc)

# Save the model as model_mobilenetv2.h5
model.save("model_mobilenetv2.h5")
print("Model saved as model_mobilenetv2.h5")