import pickle
import faiss
import streamlit as st
from sentence_transformers import SentenceTransformer
import pandas as pd

def search(query, model, embeddings, filtered_df, top_x=20):

        filtered_df_indecies_list = filtered_df.index
        filtered_embeddings = embeddings[filtered_df_indecies_list]
        
        # Load or create FAISS index
        dimension = filtered_embeddings.shape[1]  
        faiss_index = faiss.IndexFlatL2(dimension)
        faiss_index.add(filtered_embeddings)

        # Convert query to embedding
        query_embedding = model.encode([query])[0].reshape(1, -1)
        
        # Perform search
        D, I = faiss_index.search(query_embedding, k=top_x)  # Search for top x similar items
        
        # Extract the sentences corresponding to the top indices
        #print(filtered_df.columns())
        top_indecies = [i for i in I[0]]

        return filtered_df.iloc[top_indecies]