from langchain_community.vectorstores import FAISS
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain.prompts import PromptTemplate
from langchain_together import Together
from langchain.retrievers.document_compressors import EmbeddingsFilter
from langchain.retrievers import ContextualCompressionRetriever
from langchain.memory import ConversationBufferWindowMemory
from langchain.chains import ConversationalRetrievalChain
import streamlit as st
import time
st.set_page_config(page_title="Health Bot", page_icon="BG.png")

st.markdown("""
<style>
body {
    background-image: url("https://huggingface.co/spaces/Umama-at-Bluchip/Medical-Chat/blob/main/BG.png");
    background-size: cover;
}
</style>
""", unsafe_allow_html=True)


def reset_conversation():
  st.session_state.messages = []
  st.session_state.memory.clear()

if "messages" not in st.session_state:
    st.session_state.messages = []

if "memory" not in st.session_state:
    st.session_state.memory = ConversationBufferWindowMemory(k=2, memory_key="chat_history",return_messages=True) 

embeddings = HuggingFaceEmbeddings(model_name="nomic-ai/nomic-embed-text-v1",model_kwargs={"trust_remote_code":True})
db = FAISS.load_local("medchat_db", embeddings, allow_dangerous_deserialization=True)
db_retriever = db.as_retriever(search_type="similarity",search_kwargs={"k": 4})

custom_prompt_template = """
Follow these instructions clearly. This is a chat template and you are a medical, mental health, and diet plan assistant. The way you speak should be in a supportive and informative style. You are given the following pieces of information to answer the user's question correctly. You will be given context, chat history and the question. Choose only the required context based on the user's question. If the question is not related to the chat history, then don't use the history. Use chat history when required for similar related questions. While searching for the relevant information always give priority to the context given. If there are multiple medicines with the same medicine name and different strength, mention them. Always take the context related only to the question. Use your knowledge base and answer the question when the context is not related to the user's question. Utilize the provided knowledge base and search for relevant information from the context. Follow the user's question and the format closely. The answer should be abstract and concise. Understand all the context given here and generate only the answer, don't repeat the chat template in the answer. If you don't know the answer, just say that you don't know, don't try to make up your own questions and answers. Add bullet points and bold text using markdown in the required area if needed, to make it more pleasing to eyes. Most importantly, give the answers in a conversational format.

Only answer the questions of the user.

If the user greets you, do not respond with anything else apart from greeting them back.

DETECT TOPIC:
* If the question contains words like 'anxiety', 'depression', 'mood', 'therapy', 'counseling', 'stress', 'suicide' -> MENTAL HEALTH
* If the question contain words like 'diet', 'meal plan', 'nutrition', 'weight loss', 'calories', 'exercise' -> DIET 
* Otherwise -> GENERAL MEDICAL 

CONDITIONAL INSTRUCTIONS:
* MENTAL HEALTH:
    * Focus on providing supportive and empathetic responses.
    * Guide the user towards reliable mental health resources (like websites or helplines).
    * Avoid giving specific diagnoses or treatment advice. 
* DIET:
    * Ask clarifying questions about the user's goals and health conditions.
    * recommend meal plans based on user's goals and health conditions.
    * Provide general healthy eating recommendations.
    * Suggest consulting a registered dietician for personalized plans.
* GENERAL MEDICAL:
    * Utilize your medical knowledge base to provide information about conditions, symptoms, and potential treatments.


CONTEXT: {context}
CHAT HISTORY: {chat_history}
QUESTION: {question}
ANSWER:
"""

prompt = PromptTemplate(template=custom_prompt_template,
                        input_variables=['context', 'question', 'chat_history'])


llm = Together(
    model="mistralai/Mistral-7B-Instruct-v0.2",
    temperature=0.9,
    max_tokens=512,
    together_api_key="48515099b0ed4e22e56da54e50feb4adfaaa901a444b0c34bb33c66abe7b2c61"
)

qa = ConversationalRetrievalChain.from_llm(
    llm=llm,
    memory=st.session_state.memory,
    retriever=db_retriever,
    combine_docs_chain_kwargs={'prompt': prompt}
)

for message in st.session_state.messages:
    with st.chat_message(message.get("role")):
        st.write(message.get("content"))

input_prompt = st.chat_input("Chat")

if input_prompt:
    with st.chat_message("user"):
        st.write(input_prompt)

    st.session_state.messages.append({"role":"user","content":input_prompt})

    with st.chat_message("assistant"):
        with st.status("Answering...",expanded=True):
            result = qa.invoke(input=input_prompt)

            message_placeholder = st.empty()

            full_response = "**_Note: Information provided may be inaccurate. Consult a qualified doctor/medical professional for accurate advice._** \n\n\n"
        for chunk in result["answer"]:
            full_response+=chunk
            time.sleep(0.02)
            
            message_placeholder.markdown(full_response+" .")
        st.button('Reset All Chat', on_click=reset_conversation)

    st.session_state.messages.append({"role":"assistant","content":result["answer"]})