import os
import cv2
import joblib
import numpy as np
from sklearn.svm import SVC
from sklearn.model_selection import train_test_split, GridSearchCV
from sklearn.metrics import classification_report, accuracy_score
from sklearn.cluster import MiniBatchKMeans

# Path to your dataset (ensure it is extracted and organized into subfolders by class)
dataset_dir = 'natural_images'  # update with your actual path

# Initialize SIFT feature extractor
sift = cv2.SIFT_create()

# Prepare lists for storing descriptors and labels
all_descriptors = []  # will accumulate all SIFT descriptors for building the vocabulary
image_descriptors = []  # list of descriptors for each image
labels = []

# Iterate over each class folder in the dataset
for label in os.listdir(dataset_dir):
    label_folder = os.path.join(dataset_dir, label)
    if not os.path.isdir(label_folder):
        continue

    # Process each image in the folder
    for file in os.listdir(label_folder):
        file_path = os.path.join(label_folder, file)
        # Read the image in grayscale
        image = cv2.imread(file_path, cv2.IMREAD_GRAYSCALE)
        if image is None:
            continue  # Skip if image reading fails

        # Resize the image to a fixed size (optional, adjust based on your dataset)
        image = cv2.resize(image, (128, 128))

        # Extract SIFT features (keypoints and descriptors)
        keypoints, descriptors = sift.detectAndCompute(image, None)
        if descriptors is None:
            continue  # Skip images with no detected keypoints

        image_descriptors.append(descriptors)
        labels.append(label)
        # Accumulate descriptors for vocabulary creation
        all_descriptors.extend(descriptors)

# Convert the list of all descriptors to a NumPy array
all_descriptors = np.array(all_descriptors)

# Build a visual vocabulary using k-means clustering (Bag of Visual Words)
vocab_size = 50  # You can adjust the number of clusters
kmeans = MiniBatchKMeans(n_clusters=vocab_size, random_state=42)
kmeans.fit(all_descriptors)

# Create a histogram of visual words for each image
features = []
for descriptors in image_descriptors:
    # Assign each descriptor to the nearest cluster (visual word)
    words = kmeans.predict(descriptors)
    # Build a histogram with bins for each cluster
    hist, _ = np.histogram(words, bins=np.arange(vocab_size + 1))
    # Normalize the histogram (L1 normalization)
    hist = hist.astype('float32')
    hist /= (hist.sum() + 1e-7)
    features.append(hist)

# Convert lists to NumPy arrays
features = np.array(features)
labels = np.array(labels)

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

# Define parameter grid for GridSearchCV with SVM
param_grid = [
    {'kernel': ['linear'], 'C': [0.1, 1, 10, 100]},
    {'kernel': ['poly'], 'C': [0.1, 1, 10, 100], 'degree': [2, 3, 4], 'gamma': ['scale', 'auto']},
    {'kernel': ['rbf'], 'C': [0.1, 1, 10, 50, 100], 'gamma': ['scale', 'auto']},
]

svm_model = SVC(probability=True)
grid_search = GridSearchCV(svm_model, param_grid, cv=5, n_jobs=-1, verbose=3)

# Train the SVM with grid search
grid_search.fit(X_train, y_train)

print("Best parameters found:", grid_search.best_params_)
print("Best cross-validation score:", grid_search.best_score_)

# Evaluate on the test set
best_model = grid_search.best_estimator_
y_pred = best_model.predict(X_test)

print("Test Accuracy:", accuracy_score(y_test, y_pred))
print(classification_report(y_test, y_pred))

# Save the trained SVM model and the k-means vocabulary for later use
joblib.dump({'grid_search': grid_search, 'kmeans': kmeans}, 'svm_natural_images_model_with_sift.pkl')