import os
import streamlit as st
import pandas as pd
from datasets import load_from_disk
from transformers import AutoTokenizer, TFAutoModel
from constant import DRGUS_STR_LIST

if DRGUS_STR_LIST:
    Drugs = DRGUS_STR_LIST.split(',')
    Drugs = [drug.strip() for drug in Drugs] 

model_ckpt = "sentence-transformers/multi-qa-mpnet-base-dot-v1"
tokenizer = AutoTokenizer.from_pretrained(model_ckpt)
model = TFAutoModel.from_pretrained(model_ckpt, from_pt=True)

def cls_pooling(model_output):
    return model_output.last_hidden_state[:, 0]
    
def get_embeddings(text_list):
    encoded_input = tokenizer(
        text_list, padding=True, truncation=True, return_tensors="tf"
    )
    encoded_input = {k: v for k, v in encoded_input.items()}
    model_output = model(**encoded_input)
    return cls_pooling(model_output)  

embeddings_dataset = load_from_disk("data")
embeddings_dataset.add_faiss_index(column="embeddings")

def recommendations(question): 
    question_embedding = get_embeddings([question]).numpy()
    scores, samples = embeddings_dataset.get_nearest_examples(
        "embeddings", question_embedding, k=5
    )
    samples_df = pd.DataFrame.from_dict(samples)
    samples_df["scores"] = scores
    samples_df.sort_values("scores", ascending=False, inplace=True,ignore_index=True)
    return samples_df[['drugName', 'review', 'scores']]

# Create Streamlit app
st.title("Call on Doc Drug Recommendation System")

st.markdown(
    """
    <style>
    #MainMenu {visibility: hidden;}
    footer {visibility: hidden;}
    </style>
    """,
    unsafe_allow_html=True
)


# Allow users to select a default question or input their own
st.sidebar.title("Choose or Enter a Question:")
selection_type = st.sidebar.radio("Select type:", ("Select Default", "Enter Custom"))

if selection_type == "Select Default":
    selected_question = st.sidebar.selectbox("Select a question", Drugs)
    if st.sidebar.button("Show Recommendations"):
        recommendation_result = recommendations(selected_question)
        st.header(f"Top 5 Recommended Drugs for '{selected_question}':")
        st.table(recommendation_result)
else:
    default_question = "I've acne problem"
    custom_question = st.sidebar.text_input("Enter your question:", default_question)
    if st.sidebar.button("Get Recommendations"):
        if custom_question:
            custom_recommendation_result = recommendations(custom_question)
            st.header("Top 5 Recommended Drugs for Your Question:")
            st.table(custom_recommendation_result)
        else:
            st.warning("Please enter a question to get recommendations.")