import os
import numpy as np
import joblib
from PIL import Image
import streamlit as st
from sklearn.neighbors import KNeighborsClassifier
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score

# Paths and global constants
#DATASET_PATH = "path_to_your_dataset"  # Update with your dataset path
MODEL_PATH = "knn_model.pkl"
CLASS_NAMES_PATH = "class_names.pkl"
TARGET_SIZE = (64, 64)

# Helper function to load and preprocess dataset
def load_dataset(dataset_path, target_size):
    images = []
    labels = []
    class_names = sorted(os.listdir(dataset_path))  # Sort for consistent class indexing
    for idx, class_name in enumerate(class_names):
        class_path = os.path.join(dataset_path, class_name)
        if not os.path.isdir(class_path):
            continue
        for image_name in os.listdir(class_path):
            image_path = os.path.join(class_path, image_name)
            try:
                img = Image.open(image_path).convert('RGB')
                img = img.resize(target_size)
                img_array = np.array(img).flatten()  # Flatten to a single array
                images.append(img_array)
                labels.append(idx)
            except Exception as e:
                print(f"Error loading image {image_path}: {e}")
    return np.array(images), np.array(labels), class_names

# Function to train and save KNN model
def train_and_save_model(dataset_path, target_size, model_path, class_names_path):
    # Load data
    X, y, class_names = load_dataset(dataset_path, target_size)
    # Split dataset
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42, stratify=y)
    # Train KNN model
    knn = KNeighborsClassifier(n_neighbors=3)
    knn.fit(X_train, y_train)
    # Evaluate model
    y_pred = knn.predict(X_test)
    accuracy = accuracy_score(y_test, y_pred)
    print(f"KNN Model Accuracy: {accuracy * 100:.2f}%")
    # Save model and class names
    joblib.dump(knn, model_path)
    joblib.dump(class_names, class_names_path)
    print(f"Model and class names saved to {model_path} and {class_names_path}")
    return knn, class_names

# Preprocess image for prediction
def preprocess_image(image_path, target_size):
    img = Image.open(image_path).convert('RGB')
    img = img.resize(target_size)
    img_array = np.array(img).flatten()
    return img_array.reshape(1, -1)  # Add batch dimension

# Main Streamlit app
def main():
    st.title("Animal Classifier")
    st.write("Upload an image of an animal to predict its category.")

    # Load model and class names
    if not os.path.exists(MODEL_PATH) or not os.path.exists(CLASS_NAMES_PATH):
        st.error("Model or class names file not found. Please train the model first.")
        return

    knn = joblib.load(MODEL_PATH)
    class_names = joblib.load(CLASS_NAMES_PATH)

    # Upload image
    uploaded_file = st.file_uploader("Upload an image", type=["jpg", "jpeg", "png"])
    if uploaded_file:
        st.image(uploaded_file, caption="Uploaded Image", use_container_width=True)
        try:
            img_array = preprocess_image(uploaded_file, TARGET_SIZE)
            if img_array.shape[1] != knn.n_features_in_:
                st.error(f"Feature size mismatch: Model expects {knn.n_features_in_} features, but input has {img_array.shape[1]}.")
            else:
                prediction = knn.predict(img_array)
                st.success(f"The predicted animal is: {class_names[prediction[0]]}")
        except Exception as e:
            st.error(f"Error processing image: {e}")

# Train and save model (Uncomment if running for the first time)
# train_and_save_model(DATASET_PATH, TARGET_SIZE, MODEL_PATH, CLASS_NAMES_PATH)

# Run Streamlit app
if __name__ == "__main__":
    main()