import glob
import cv2
import numpy as np
import gradio as gr
import matplotlib.pyplot as plt
from sklearn.model_selection import KFold, GridSearchCV
from sklearn.neighbors import KNeighborsClassifier
from sklearn.metrics import (
    accuracy_score,
    classification_report,
    confusion_matrix,
    precision_score,
    recall_score,
    f1_score,
)
from skimage.feature import graycomatrix, graycoprops, local_binary_pattern

# Define directories for grass and wood images
grass_dir = "images/Grass/Train_Grass"
wood_dir = "images/Wood/Train_wood"

# Constants for Local Binary Pattern (LBP) and Gray Level Co-occurrence Matrix (GLCM)
RADIUS = 1
N_POINTS = 12 * RADIUS
TARGET_SIZE = (30, 30)
distances = [1]
angles = [0]


def load_and_convert_images(directory):
    """Load images from a specified directory using glob and convert them to grayscale.

    Args:
        directory (str): The path to the image directory.

    Returns:
        list: A list of resized grayscale images.
    """
    dataset = []
    for img_path in glob.glob(f"{directory}/*.*"):
        if img_path.endswith((".jpg", ".png", ".jpeg")):
            img = cv2.imread(img_path)  # Read the image
            if img is not None:
                # Convert to grayscale and resize the image
                gray_image = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
                resized_image = cv2.resize(
                    gray_image, TARGET_SIZE, interpolation=cv2.INTER_AREA
                )
                dataset.append(resized_image)
    return dataset


# Load datasets using glob
grass_dataset = load_and_convert_images(grass_dir)
wood_dataset = load_and_convert_images(wood_dir)


def calc_glcm_features(images):
    """Calculate GLCM features for a list of images.

    Args:
        images (list): A list of grayscale images.

    Returns:
        list: A list of GLCM features for each image.
    """
    features = []
    for img in images:
        glcm = graycomatrix(img, distances, angles, symmetric=True, normed=True)

        # Calculate GLCM properties
        contrast = graycoprops(glcm, "contrast")[0, 0]
        dissimilarity = graycoprops(glcm, "dissimilarity")[0, 0]
        homogeneity = graycoprops(glcm, "homogeneity")[0, 0]
        energy = graycoprops(glcm, "energy")[0, 0]
        correlation = graycoprops(glcm, "correlation")[0, 0]

        features.append([contrast, dissimilarity, homogeneity, energy, correlation])
    return features


# Calculate GLCM features for grass and wood datasets
grass_glcm_features = calc_glcm_features(grass_dataset)
wood_glcm_features = calc_glcm_features(wood_dataset)

# Print the sizes of the GLCM feature arrays
print(f"Size of GLCM features for grass dataset: {len(grass_glcm_features)}")
print(f"Size of GLCM features for wood dataset: {len(wood_glcm_features)}")


def extract_lbp_features(images):
    """Extract LBP features from a list of images.

    Args:
        images (list): A list of grayscale images.

    Returns:
        list: A list of LBP histograms for each image.
    """
    lbp_features = []
    for image in images:
        lbp = local_binary_pattern(image, N_POINTS, RADIUS, method="uniform")
        n_bins = int(lbp.max() + 1)
        lbp_hist, _ = np.histogram(lbp, bins=n_bins, range=(0, n_bins), density=True)
        lbp_features.append(lbp_hist)
    return lbp_features


# Extract LBP features for grass and wood datasets
grass_lbp_features = extract_lbp_features(grass_dataset)
wood_lbp_features = extract_lbp_features(wood_dataset)

# Create labels (0 for grass, 1 for wood)
grass_labels = [0] * len(grass_dataset)  # Label all grass images as 0
wood_labels = [1] * len(wood_dataset)  # Label all wood images as 1

# Combine features and labels for GLCM classifier
glcm_features = np.array(grass_glcm_features + wood_glcm_features)
glcm_labels = grass_labels + wood_labels

# Prepare labels and features for LBP classifier
lbp_features = np.array(grass_lbp_features + wood_lbp_features)
lbp_labels = grass_labels + wood_labels

# Number of images
num_grass = len(grass_dataset)
num_wood = len(wood_dataset)

# Define labels for classification
y = np.array([0] * num_grass + [1] * num_wood)

# Define KFold cross-validation
k = 5
kf = KFold(n_splits=k, shuffle=True, random_state=42)

# Store results for GLCM and LBP classifiers
glcm_metrics = {"accuracy": [], "precision": [], "recall": [], "f1_score": []}
lbp_metrics = {"accuracy": [], "precision": [], "recall": [], "f1_score": []}
y_true_glcm, y_true_lbp = [], []
y_pred_glcm, y_pred_lbp = [], []

# Parameter tuning using GridSearchCV for KNN classifier
param_grid = {"n_neighbors": [3, 5, 7], "p": [1, 2]}

# GLCM Classifier Training and Evaluation
glcm_knn = KNeighborsClassifier()
glcm_grid_search = GridSearchCV(glcm_knn, param_grid, cv=kf)
glcm_grid_search.fit(glcm_features, y)

# Perform cross-validation and evaluate GLCM classifier
for train_index, test_index in kf.split(glcm_features):
    x_train, x_test = glcm_features[train_index], glcm_features[test_index]
    y_train, y_test = y[train_index], y[test_index]

    glcm_classifier = KNeighborsClassifier(
        n_neighbors=glcm_grid_search.best_params_["n_neighbors"]
    )
    glcm_classifier.fit(x_train, y_train)
    y_pred = glcm_classifier.predict(x_test)

    # Collect true and predicted labels
    y_true_glcm.extend(y_test)
    y_pred_glcm.extend(y_pred)

    # Calculate and store metrics
    accuracy = accuracy_score(y_test, y_pred)
    precision = precision_score(y_test, y_pred, average="macro")
    recall = recall_score(y_test, y_pred, average="macro")
    f1 = f1_score(y_test, y_pred, average="macro")

    glcm_metrics["accuracy"].append(accuracy)
    glcm_metrics["precision"].append(precision)
    glcm_metrics["recall"].append(recall)
    glcm_metrics["f1_score"].append(f1)

    # Print metrics for this fold
    print(
        f"GLCM Fold Metrics: Accuracy={accuracy:.2f}, Precision={precision:.2f}, Recall={recall:.2f}, F1-Score={f1:.2f}"
    )

# Print average metrics for GLCM classifier
print("\nAverage GLCM Classifier Metrics:")
for metric in glcm_metrics:
    avg_metric = np.mean(glcm_metrics[metric])
    print(f"{metric.capitalize()}: {avg_metric:.2f}")

# Confusion Matrix for GLCM
glcm_conf_matrix = confusion_matrix(y_true_glcm, y_pred_glcm)
print("\nConfusion Matrix for GLCM Classifier:")
print(glcm_conf_matrix)

# Classification Report for GLCM
print("\nClassification Report for GLCM Classifier:")
print(classification_report(y_true_glcm, y_pred_glcm))

# LBP Classifier Training and Evaluation
lbp_knn = KNeighborsClassifier()
lbp_grid_search = GridSearchCV(lbp_knn, param_grid, cv=kf)
lbp_grid_search.fit(lbp_features, y)

# Perform cross-validation and evaluate LBP classifier
for train_index, test_index in kf.split(lbp_features):
    x_train, x_test = lbp_features[train_index], lbp_features[test_index]
    y_train, y_test = y[train_index], y[test_index]

    lbp_classifier = KNeighborsClassifier(
        n_neighbors=lbp_grid_search.best_params_["n_neighbors"]
    )
    lbp_classifier.fit(x_train, y_train)
    y_pred = lbp_classifier.predict(x_test)

    # Collect true and predicted labels
    y_true_lbp.extend(y_test)
    y_pred_lbp.extend(y_pred)

    # Calculate and store metrics
    accuracy = accuracy_score(y_test, y_pred)
    precision = precision_score(y_test, y_pred, average="macro")
    recall = recall_score(y_test, y_pred, average="macro")
    f1 = f1_score(y_test, y_pred, average="macro")

    lbp_metrics["accuracy"].append(accuracy)
    lbp_metrics["precision"].append(precision)
    lbp_metrics["recall"].append(recall)
    lbp_metrics["f1_score"].append(f1)

    # Print metrics for this fold
    print(
        f"LBP Fold Metrics: Accuracy={accuracy:.2f}, Precision={precision:.2f}, Recall={recall:.2f}, F1-Score={f1:.2f}"
    )

# Print average metrics for LBP classifier
print("\nAverage LBP Classifier Metrics:")
for metric in lbp_metrics:
    avg_metric = np.mean(lbp_metrics[metric])
    print(f"{metric.capitalize()}: {avg_metric:.2f}")

# Confusion Matrix for LBP
lbp_conf_matrix = confusion_matrix(y_true_lbp, y_pred_lbp)
print("\nConfusion Matrix for LBP Classifier:")
print(lbp_conf_matrix)

# Classification Report for LBP
print("\nClassification Report for LBP Classifier:")
print(classification_report(y_true_lbp, y_pred_lbp))


def classify_texture(image, method):
    """Classify the texture of the uploaded image as grass or wood using the selected method.

    Args:
        image (numpy.ndarray): The uploaded image to classify.
        method (str): The feature extraction method ('GLCM' or 'LBP').

    Returns:
        Tuple[str, numpy.ndarray]: The classification result ('Grass' or 'Wood') and the highlighted image.
    """
    # Pre-process the image
    gray_image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    resized_image = cv2.resize(gray_image, TARGET_SIZE, interpolation=cv2.INTER_AREA)

    # Extract features based on the selected method
    if method == "GLCM":
        feature = calc_glcm_features([resized_image])
        prediction = glcm_classifier.predict(feature)
    elif method == "LBP":
        feature = extract_lbp_features([resized_image])
        prediction = lbp_classifier.predict(feature)
    else:
        raise ValueError("The method is not recognized")

    # Classify the prediction
    result = "Grass" if prediction == 0 else "Wood"

    # Highlight the image based on the classification
    if result == "Grass":
        highlighted_image = cv2.cvtColor(
            image, cv2.COLOR_BGR2RGB
        )  # Convert to RGB for displaying
        highlighted_image[:] = [0, 255, 0]  # Highlight in green
    else:
        highlighted_image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        highlighted_image[:] = [165, 42, 42]  # Highlight in brown

    return result, highlighted_image


# Gradio interface setup with a dropdown for method selection
iface = gr.Interface(
    fn=classify_texture,
    inputs=[
        gr.Image(type="numpy", label="Upload Image"),
        gr.Dropdown(choices=["Choose Here", "GLCM", "LBP"], label="Select Feature Extraction Method"),
    ],
    outputs=["text", gr.Image(type="numpy", label="Highlighted Image")],
    title="Texture Classification",
    description="Upload an image of grass or wood to classify the texture. Select GLCM or LBP as the method.",
)

# Launch Gradio interface
iface.launch()