import streamlit as st
import os
import pickle as pkl
import numpy as np
from sklearn.neighbors import NearestNeighbors
from tensorflow.keras.applications.resnet50 import ResNet50, preprocess_input
from tensorflow.keras.preprocessing import image
from tensorflow.keras.layers import GlobalMaxPool2D
import tensorflow as tf

# ---------------------------------------
# Page title and description
st.set_page_config(page_title="Fashion Product Recommendation System", layout="wide")
st.title("Fashion Product Recommendation System - No Images")
st.write(
    "This system will show the 5 most similar products to the uploaded image. "
    "Note: Images are not available in the Space environment, so only product IDs are displayed."
)

# ---------------------------------------
# Script directory
BASE_DIR = os.path.dirname(__file__)

# Pickle file paths
features_path = os.path.join(BASE_DIR, "Images_features.pkl")
filenames_path = os.path.join(BASE_DIR, "filenames.pkl")

# ---------------------------------------
# Load pickle files
image_features = pkl.load(open(features_path, "rb"))
filenames = pkl.load(open(filenames_path, "rb"))

# Keep only base names for IDs
filenames = [os.path.basename(f) for f in filenames]

# ---------------------------------------
# Create k-NN model
neighbors = NearestNeighbors(n_neighbors=6, algorithm='brute', metric='euclidean')
neighbors.fit(image_features)

# ---------------------------------------
# Feature extraction model
base_model = ResNet50(weights='imagenet', include_top=False, input_shape=(224,224,3))
base_model.trainable = False
model = tf.keras.models.Sequential([base_model, GlobalMaxPool2D()])

# ---------------------------------------
# Function to extract features from image
def extract_features_from_image_file(img_file, model):
    img = image.load_img(img_file, target_size=(224,224))
    arr = image.img_to_array(img)
    arr = np.expand_dims(arr, axis=0)
    arr = preprocess_input(arr)
    vec = model.predict(arr, verbose=0).flatten()
    vec /= np.linalg.norm(vec) + 1e-10  # Normalize
    return vec

# ---------------------------------------
# File uploader
uploaded_file = st.file_uploader(
    "Please upload a product image (jpg/png):",
    type=["jpg", "jpeg", "png"]
)

if uploaded_file is not None:
    st.write("Extracting features...")
    features = extract_features_from_image_file(uploaded_file, model)
    
    # Find nearest neighbors
    dists, idxs = neighbors.kneighbors([features])
    
    st.subheader("Top 5 Similar Products (IDs only)")
    for i, idx in enumerate(idxs[0][1:]):  # skip first index (itself)
        st.write(f"{i+1}: {filenames[idx]}")
else:
    st.info("Please upload an image to use the system.")