# streamlit_app.py
import streamlit as st
import pandas as pd
import torch
from sentence_transformers import SentenceTransformer, util
import pickle
import numpy as np

import os
import importlib

#Load sentences & embeddings from disc
with open('clinical_inno_embeddings_masterid_paraphrase-multilingual-mpnet-base-v2.pkl', "rb") as fIn:
    stored_data = pickle.load(fIn)
    stored_masterid = stored_data['pro_master_id']
    stored_products = stored_data['products']
    stored_embeddings = stored_data['embeddings']

# Initialize the SentenceTransformer model
embedder = SentenceTransformer('sentence-transformers/paraphrase-multilingual-mpnet-base-v2')

def get_similar_products(query, products, mean_embeddings_tensor, top_k=10):
    query_embedding = embedder.encode(query, convert_to_tensor=True)
    cos_scores = util.cos_sim(query_embedding, stored_embeddings)[0]
    top_results = torch.topk(cos_scores, k=top_k)
    
    similar_products = [(products[idx.item()], score.item()) for score, idx in zip(top_results[0], top_results[1])]
    return similar_products


# Streamlit UI
st.title("Product Similarity Finder")

# User input
user_input = st.text_input("Enter a product name or description:")

# Search button
if st.button('Search'):
    if user_input:
        # Get and display similar products
        results = get_similar_products(user_input, stored_products, stored_embeddings)
        
        # Convert results to a DataFrame for nicer display
        results_df = pd.DataFrame(results, columns=['Product', 'Score'])
        
        # Use Streamlit's dataframe function to display results in a table with default formatting
        st.dataframe(results_df.style.format({'Score': '{:.4f}'}))
    else:
        st.write("Please enter a product name or description to search.")