import streamlit as st
from transformers import pipeline
import faiss
import numpy as np
import os

# Initialize models and retrieval pipeline
@st.cache_resource
def load_model():
    # Load the retrieval and generation models
    retrieval_model = pipeline('question-answering', model="bert-large-uncased-whole-word-masking-finetuned-squad")
    generation_model = pipeline('text-generation', model="gpt-2")
    return retrieval_model, generation_model

retrieval_model, generation_model = load_model()

# Function to search relevant medical documents using FAISS
def search_medical_documents(query, faiss_index):
    # Convert the query to a vector (embedding) using a model (e.g., BERT or any embedding model)
    query_vector = retrieval_model.tokenizer(query, return_tensors="pt", padding=True, truncation=True)
    query_embedding = retrieval_model.model(**query_vector).last_hidden_state.mean(dim=1).detach().numpy()

    # Perform search in FAISS index
    D, I = faiss_index.search(query_embedding, k=5)  # Top 5 results
    retrieved_docs = [faiss_index.reconstruct(i) for i in I[0]]  # Retrieve documents based on indices
    return retrieved_docs

# Build FAISS index (you should have your own documents here)
def build_faiss_index(documents):
    embeddings = [retrieval_model.tokenizer(doc, return_tensors="pt", padding=True, truncation=True) for doc in documents]
    embeddings = [retrieval_model.model(**embed).last_hidden_state.mean(dim=1).detach().numpy() for embed in embeddings]

    dim = len(embeddings[0])  # Dimensionality of the embedding
    faiss_index = faiss.IndexFlatL2(dim)
    faiss_index.add(np.array(embeddings))  # Add embeddings to the index
    return faiss_index

# Set up the Streamlit app interface
def chatbot_interface():
    st.title("AI-Powered Healthcare Chatbot")
    st.write("Ask any health-related questions or upload medical records for analysis.")

    user_input = st.text_input("Enter your query here:")

    if user_input:
        # Retrieve relevant documents using the FAISS index (after indexing documents)
        retrieved_docs = search_medical_documents(user_input, faiss_index)

        context = " ".join(retrieved_docs)

        response = generation_model(f"Context: {context} \nQuery: {user_input}\nResponse:")
        
        st.write("Response: ")
        st.write(response[0]['generated_text'])

    # File uploader for medical records
    uploaded_file = st.file_uploader("Upload medical documents (PDF, DOCX, etc.):", type=["pdf", "docx"])
    if uploaded_file is not None:
        extracted_text = extract_text_from_file(uploaded_file)
        st.write("Extracted Text: ")
        st.write(extracted_text)

        response = generation_model(f"Extracted Text: {extracted_text}\nAnswer the following question: {user_input}")
        st.write("Generated Response:")
        st.write(response[0]['generated_text'])

# Function to extract text from uploaded medical files
def extract_text_from_file(uploaded_file):
    file_name = uploaded_file.name
    if file_name.endswith(".pdf"):
        return extract_text_from_pdf(uploaded_file)
    elif file_name.endswith(".docx"):
        return extract_text_from_docx(uploaded_file)
    else:
        return "File type not supported."

if __name__ == "__main__":
    chatbot_interface()