import streamlit as st
import torch
import torch.nn as nn
import torchvision.models as models
import torchvision.transforms as transforms
from sklearn.neighbors import NearestNeighbors
from PIL import Image

# Define the path to the directory containing the images
IMAGE_DIR = "lfw"

# Define the path to the ResNet50 model checkpoint
MODEL_CHECKPOINT = "resnet50.pth"

# Define the number of nearest neighbors to retrieve
NUM_NEIGHBORS = 10

# Load the pretrained ResNet50 model
model = models.resnet50(pretrained=True)

# Remove the last layer (the classification layer) from the model
model = nn.Sequential(*list(model.children())[:-1])

# Load the saved ResNet50 model checkpoint
model.load_state_dict(torch.load(MODEL_CHECKPOINT, map_location=torch.device('cpu')))

# Set the model to evaluation mode
model.eval()

# Define a preprocessing transform to resize and normalize the images
transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))
])

# Load the image filenames and their corresponding feature vectors
image_filenames = []
features = []
with open("features.txt", "r") as f:
    for line in f:
        parts = line.strip().split(",")
        image_filenames.append(parts[0])
        feature = [float(x) for x in parts[1:]]
        features.append(feature)
features = torch.tensor(features)

# Create a nearest neighbor model and fit it to the feature vectors
model = NearestNeighbors(n_neighbors=NUM_NEIGHBORS, metric='euclidean')
model.fit(features)

# Define a function to find the 10 most similar images to a query image
def find_similar_images(query_image_path):
    # Load and preprocess the query image
    query_image = Image.open(query_image_path)
    query_image = transform(query_image)

    # Extract the feature vector from the query image
    query_feature = model(torch.unsqueeze(query_image, 0))
    query_feature = query_feature.reshape(query_feature.shape[0], -1).detach().numpy()

    # Find the indices of the 10 most similar images
    distances, indices = model.kneighbors(query_feature)

    # Return the paths to the 10 most similar images
    similar_image_paths = [image_filenames[i] for i in indices[0]]
    return similar_image_paths

# Define the Streamlit app
def app():
    # Set the page title
    st.set_page_config(page_title="Similarity Search App", page_icon=":mag_right:")

    # Define the sidebar
    st.sidebar.title("Similarity Search")
    query_image_method = st.sidebar.radio("Select method:", ("Select image", "Upload image"))

    # Define the main content
    st.title("Similarity Search App")

    if query_image_method == "Select image":
        # List all the available images
        image_files = [f"{IMAGE_DIR}/{name}" for name in os.listdir(IMAGE_DIR)]
        selected_image = st.selectbox("Select an image", image_files)

        # Display the selected image
        st.image(selected_image, caption="Selected Image", use_column_width=True)

        # Find the most similar images
        similar_images = find_similar_images(selected_image)

        # Display the most similar images
        st.subheader("Similar Images")
        for i, image_path in enumerate(similar_images):
            image = Image.open(image_path)
            st.image(image, caption=f"Rank {i+1}", width=150)