import os
import numpy as np
import pickle
import tensorflow as tf
from PIL import Image
from tensorflow.keras.preprocessing import image
from tensorflow.keras.layers import GlobalMaxPooling2D
from tensorflow.keras.applications.resnet50 import ResNet50, preprocess_input
from sklearn.neighbors import NearestNeighbors
from numpy.linalg import norm
import gradio as gr

# Function to load embeddings and filenames
def load_data():
    # Load embeddings and filenames from the same directory as main.py
    embeddings_path = os.path.join(os.path.dirname(__file__), 'embeddings.pkl')
    filenames_path = os.path.join(os.path.dirname(__file__), 'filenames.pkl')
    
    feature_list = np.array(pickle.load(open(embeddings_path, 'rb')))
    filenames = pickle.load(open(filenames_path, 'rb'))
    
    return feature_list, filenames

# Load the feature list and filenames
feature_list, filenames = load_data()

# Load the pre-trained model
model = ResNet50(weights='imagenet', include_top=False, input_shape=(224, 224, 3))
model.trainable = False
model = tf.keras.Sequential([
    model,
    GlobalMaxPooling2D()
])

def feature_extraction(img_path, model):
    img = image.load_img(img_path, target_size=(224, 224))
    img_array = image.img_to_array(img)
    expanded_img_array = np.expand_dims(img_array, axis=0)
    preprocessed_img = preprocess_input(expanded_img_array)
    result = model.predict(preprocessed_img).flatten()
    normalized_result = result / norm(result)
    return normalized_result

def recommend(features, feature_list):
    neighbors = NearestNeighbors(n_neighbors=6, algorithm='brute', metric='euclidean')
    neighbors.fit(feature_list)
    distances, indices = neighbors.kneighbors([features])
    return indices

def resize_image(img, size):
    img = img.resize(size, Image.LANCZOS)
    return img

def predict(input_image):
    img = input_image.convert("RGB")  # Ensure image is in RGB mode
    img_array = np.array(img)
    expanded_img_array = np.expand_dims(img_array, axis=0)
    preprocessed_img = preprocess_input(expanded_img_array)
    features = model.predict(preprocessed_img).flatten()
    normalized_features = features / norm(features)
    indices = recommend(normalized_features, feature_list)
    result_images = [filenames[idx] for idx in indices[0]]
    resized_images = [resize_image(Image.open(img_path), (224, 224)) for img_path in result_images]
    return resized_images

iface = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil", label="Upload an image"),
    outputs=[gr.Image(type="pil", label=f"Recommendation {i+1}") for i in range(5)],
    title="Amazon Lens",
    description="Upload an image of clothing, and the system will recommend similar items."
)

if __name__ == "__main__":
    iface.launch()