# Optimise test data pipeline
AUTOTUNE = tf.data.AUTOTUNE
test_ds = test_ds.cache().prefetch(buffer_size=AUTOTUNE)
# Function to create model
def create_model(learning_rate):
    data_augmentation = keras.Sequential([
        layers.RandomFlip("horizontal"),
        layers.RandomRotation(0.1),
        layers.RandomZoom(0.1),
    ])

    model = keras.Sequential([
        layers.Input(shape=(img_height, img_width, 3)),
        data_augmentation,
        layers.Rescaling(1./255),

        layers.Conv2D(32, 3, activation="relu", padding="same"),
        layers.MaxPooling2D(),

        layers.Conv2D(64, 3, activation="relu", padding="same"),
        layers.MaxPooling2D(),

        layers.Conv2D(128, 3, activation="relu", padding="same"),
        layers.MaxPooling2D(),

        layers.Flatten(),
        layers.Dense(128, activation="relu"),
        layers.Dropout(0.3),
        layers.Dense(num_classes, activation="softmax")
    ])

    model.compile(
        optimizer=keras.optimizers.Adam(learning_rate=learning_rate),
        loss="sparse_categorical_crossentropy",
        metrics=["accuracy"]
    )

    return model

	# Hyperparameter tuning - Train multiple models
import time

results = []

for idx, params in enumerate(hyperparameters):
    print(f"Testing Configuration {idx+1}: {params}")

    # Load training and validation data with current batch size
    train_ds = tf.keras.utils.image_dataset_from_directory(
        train_dir,
        validation_split=0.2,
        subset="training",
        seed=seed,
        image_size=(img_height, img_width),
        batch_size=params['batch_size']
    )

    val_ds = tf.keras.utils.image_dataset_from_directory(
        train_dir,
        validation_split=0.2,
        subset="validation",
        seed=seed,
        image_size=(img_height, img_width),
        batch_size=params['batch_size']
    )

    # Optimize data pipeline
    train_ds = train_ds.cache().shuffle(1000).prefetch(buffer_size=AUTOTUNE)
    val_ds = val_ds.cache().prefetch(buffer_size=AUTOTUNE)

    # Create and train model
    model = create_model(params['learning_rate'])

    start_time = time.time()

    history = model.fit(
        train_ds,
        validation_data=val_ds,
        epochs=params['epochs'],
        verbose=1
    )

    training_time = time.time() - start_time

    # Evaluate on test set
    test_loss, test_acc = model.evaluate(test_ds, verbose=0)

    # Store results
    results.append({
        'params': params,
        'history': history.history,
        'test_accuracy': test_acc,
        'test_loss': test_loss,
        'training_time': training_time,
        'model': model
    })

    print(f"\nTest Accuracy: {test_acc*100:.2f}%")
    print(f"Test Loss: {test_loss:.4f}")
    print(f"Training Time: {training_time:.2f} seconds")
	# Compare hyperparameter results

print("Hyperparameter Comparison")

for idx, result in enumerate(results):
    print(f"\nConfiguration {idx+1}: {result['params']}")
    print(f"Test Accuracy: {result['test_accuracy']*100:.2f}%")
    print(f"Test Loss: {result['test_loss']:.4f}")
    print(f"Training Time: {result['training_time']:.2f} seconds")

best_idx = np.argmax([r['test_accuracy'] for r in results])
print(f"\nBest Configuration: {best_idx+1}")
print(f"Best Test Accuracy: {results[best_idx]['test_accuracy']*100:.2f}%")

best_model = results[best_idx]['model']
# Plot training history for all configurations
num_configs = len(results)
fig, axes = plt.subplots(num_configs, 2, figsize=(15, 5*num_configs))

if num_configs == 1:
    axes = axes.reshape(1, -1)

for idx, result in enumerate(results):
    history = result['history']
    epochs_range = range(1, len(history['accuracy']) + 1)

    # Plot accuracy
    axes[idx, 0].plot(epochs_range, history['accuracy'], label='Train Accuracy')
    axes[idx, 0].plot(epochs_range, history['val_accuracy'], label='Val Accuracy')
    axes[idx, 0].set_xlabel('Epoch')
    axes[idx, 0].set_ylabel('Accuracy')
    axes[idx, 0].set_title(f'Config {idx+1}: Accuracy - {result["params"]}')
    axes[idx, 0].legend()
    axes[idx, 0].grid(True)

    # Plot loss
    axes[idx, 1].plot(epochs_range, history['loss'], label='Train Loss')
    axes[idx, 1].plot(epochs_range, history['val_loss'], label='Val Loss')
    axes[idx, 1].set_xlabel('Epoch')
    axes[idx, 1].set_ylabel('Loss')
    axes[idx, 1].set_title(f'Config {idx+1}: Loss - {result["params"]}')
    axes[idx, 1].legend()
    axes[idx, 1].grid(True)

plt.tight_layout()
plt.show()
# Evaluate best model on test set
test_loss, test_acc = best_model.evaluate(test_ds)
print(f"Best Model Test Accuracy: {test_acc*100:.2f}%")
print(f"Best Model Test Loss: {test_loss:.4f}")
# Confusion Matrix
y_true = np.concatenate([labels.numpy() for _, labels in test_ds], axis=0)
y_pred_probs = best_model.predict(test_ds)
y_pred = np.argmax(y_pred_probs, axis=1)

cm = confusion_matrix(y_true, y_pred)
disp = ConfusionMatrixDisplay(confusion_matrix=cm, display_labels=class_names)
fig, ax = plt.subplots(figsize=(10, 8))
disp.plot(ax=ax, cmap="Blues", xticks_rotation=45)
plt.title("Confusion Matrix")
plt.show()
# Classification Report
print("\nClassification Report:")
print(classification_report(y_true, y_pred, target_names=class_names))
# Model Prediction on Random 6 Images
import random

pred_files = sorted([
    os.path.join(pred_dir, f)
    for f in os.listdir(pred_dir)
    if f.lower().endswith((".jpg", ".jpeg", ".png"))
])

sample_files = random.sample(pred_files, 6)

fig, axes = plt.subplots(2, 3, figsize=(15, 10))
axes = axes.flatten()

for i, file_path in enumerate(sample_files):
    img = tf.keras.utils.load_img(file_path, target_size=(img_height, img_width))
    img_array = tf.keras.utils.img_to_array(img)
    img_array_exp = tf.expand_dims(img_array, 0)

    prediction = best_model.predict(img_array_exp, verbose=0)
    predicted_class = class_names[np.argmax(prediction)]

    axes[i].imshow(img)
    axes[i].set_title(f'Predicted: {predicted_class}')
    axes[i].axis('off')

plt.tight_layout()
plt.show()