import pandas as pd


df = pd.read_csv('./drugs_side_effects_drugs_com.csv')

df = df[['drug_name', 'medical_condition', 'side_effects']]
df.dropna(inplace=True)

context_data = []
for i in range(10):
  context = ""
  for j in range(3):
    context += df.columns[j]
    context += ": "
    context += str(df.iloc[i][j])
    context += " "
  context_data.append(context)

import os

# Get the secret key from the environment
groq_key = os.environ.get('gloq_key')

## LLM used for RAG
from langchain_groq import ChatGroq

llm = ChatGroq(model="llama-3.1-70b-versatile",api_key=groq_key)

## Embedding model!
from langchain_huggingface import HuggingFaceEmbeddings
embed_model = HuggingFaceEmbeddings(model_name="mixedbread-ai/mxbai-embed-large-v1")

# create vector store!
from langchain_chroma import Chroma

vectorstore = Chroma(
    collection_name="medical_dataset_store",
    embedding_function=embed_model,
    persist_directory="./",
)

# add data to vector nstore
vectorstore.add_texts(context_data)

retriever = vectorstore.as_retriever()

from langchain_core.prompts import PromptTemplate

template = ("""You are a pharmacist and medical expert. 
    Use the provided context to answer the question.
    If the question is related to medical condition, drug name 
    and side effects that are not in the context, look online and answer them.
    If you don't know the answer, say so. Explain your answer in detail.
    Do not discuss the context in your response; just provide the answer directly.

    Context: {context}

    Question: {question}

    Answer:""")

rag_prompt = PromptTemplate.from_template(template)

from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables import RunnablePassthrough

rag_chain = (
    {"context": retriever, "question": RunnablePassthrough()}
    | rag_prompt
    | llm
    | StrOutputParser()
)


import gradio as gr

# Function to stream responses
def rag_memory_stream(text):
    partial_text = ""
    for new_text in rag_chain.stream(text):  # Assuming rag_chain is pre-defined
        partial_text += new_text
        yield partial_text

# Title and description for the app
title = "AI Medical Assistant for Drug Information and Side Effects"
description = """
<div class="description">
This AI-powered chatbot is designed to provide reliable information about drugs, their side effects, and related medical conditions.
It utilizes the Groq API and LangChain to deliver real-time, accurate responses.

Ask questions like:
<ul>
    <li>What are the side effects of taking aspirin daily?</li>
    <li>What is the recommended treatment for a common cold?</li>
    <li>What is the disease for constant fatigue and muscle weakness?</li>
    <li>What are the symptoms of diabetes?</li>
    <li>How can hypertension be managed?</li>
</ul>

<strong>Disclaimer:</strong> This chatbot is for informational purposes only and is not a substitute for professional medical advice.
</div>
"""

# Customizing Gradio interface for a better look

# HTML for custom styling
custom_css = """
body {
    background-color: #f9f9f9;
    font-family: Arial, sans-serif;
    margin: 0;
    padding: 0;
}

#interface-container {
    max-width: 800px;
    margin: 50px auto;
    padding: 20px;
    background: linear-gradient(145deg, #000000, #f0f0f0);
    box-shadow: 0 4px 6px rgba(0, 0, 0, 0.1);
    border-radius: 10px;
    border: 1px solid #e0e0e0;
}

h1 {
    text-align: center;
    color: #333;
}

.description {
    text-align: justify;
    color: #000;
    font-size: 1rem;
    margin-bottom: 20px;
}

footer {
    text-align: center;
    color: #777;
    margin-top: 30px;
    font-size: 0.9rem;
}
"""

# Customizing Gradio interface with additional CSS and content
demo = gr.Interface(
    fn=rag_memory_stream,
    inputs=gr.Textbox(
        lines=5, 
        placeholder="Type your medical question here...", 
        label="Your Medical Question"
    ),
    outputs=gr.Textbox(
        lines=15,  # Reduced line count for better layout
        label="AI Response"
    ),
    title=title,
    description=description,
    css=custom_css,
    allow_flagging="never"
)

if __name__ == "__main__":
    demo.launch()