import os
import streamlit as st
from PIL import Image
import numpy as np
import pickle
import tensorflow as tf
from tensorflow.keras.preprocessing import image
from tensorflow.keras.layers import GlobalMaxPooling2D
from sklearn.neighbors import NearestNeighbors
from numpy.linalg import norm
from classification_models.tfkeras import Classifiers

# Function to gather file paths
def get_image_paths():
    filenames = []
    folder_names = sorted(['images1', 'images2', 'images3', 'images4', 'images5'])
    for folder_name in folder_names:
        for file in os.listdir(folder_name):
            if not file.endswith('.lnk'):
                filenames.append(os.path.join(folder_name, file))
    return filenames

# Gather image paths
filenames = get_image_paths()

# Load precomputed features (ensure this file is available)
feature_list = pickle.load(open('feature_list.pkl', 'rb'))
feature_list = np.array(feature_list)

# Get the ResNeXt model
ResNeXt50, preprocess_input = Classifiers.get('resnext50')
model = ResNeXt50(include_top=False, input_shape=(224, 224, 3), weights='imagenet')
model.trainable = False
model = tf.keras.Sequential([model, GlobalMaxPooling2D()])

# App title
st.title('G Fashion')

def save_uploaded_file(uploaded_file):
    try:
        if not os.path.exists('uploads'):
            os.makedirs('uploads')
        with open(os.path.join('uploads', uploaded_file.name), 'wb') as f:
            f.write(uploaded_file.getbuffer())
        return 1
    except Exception as e:
        st.error(f"Error in file upload: {e}")
        return 0

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='cosine')
    neighbors.fit(feature_list)
    distances, indices = neighbors.kneighbors([features])
    return indices

# File upload
uploaded_file = st.file_uploader("Choose an image")
if uploaded_file is not None:
    if save_uploaded_file(uploaded_file):
        display_image = Image.open(uploaded_file)
        st.image(display_image)
        features = feature_extraction(os.path.join("uploads", uploaded_file.name), model)
        indices = recommend(features, feature_list)
        col1, col2, col3, col4, col5 = st.columns(5)

        st.write(f"File path: {filenames[indices[0][0]]}")

        with col1:
            st.image(r'{}'.format(filenames[indices[0][0]]))
        with col2:
            st.image(r'{}'.format(filenames[indices[0][1]]))
        with col3:
            st.image(r'{}'.format(filenames[indices[0][2]]))
        with col4:
            st.image(r'{}'.format(filenames[indices[0][3]]))
        with col5:
            st.image(r'{}'.format(filenames[indices[0][4]]))
    else:
        st.header("Some error occured in file upload")