import streamlit as st
import pdfplumber
from transformers import pipeline

# Load AI Models
summarizer = pipeline("summarization", model="facebook/bart-large-cnn")
qa_model = pipeline("question-answering", model="deepset/roberta-base-squad2")

# Function to extract text from PDF
def extract_text_from_pdf(pdf):
    with pdfplumber.open(pdf) as pdf_file:
        text = ""
        for page in pdf_file.pages:
            text += page.extract_text() + "\n"
    return text

# Function to summarize text
def summarize_text(text):
    # Limit input to 1024 characters (Bart-large-cnn model limit)
    max_input_length = 1024
    text = text[:max_input_length]  # Truncate text to avoid errors
    
    summary = summarizer(text, max_length=200, min_length=50, do_sample=False)
    return summary[0]['summary_text']


# Function for Q&A
def answer_question(context, question):
    response = qa_model(question=question, context=context)
    return response['answer']

# Streamlit UI
st.set_page_config(page_title="MedGen-AI", page_icon="🩺", layout="wide")
st.title("🩺 MedGen-AI: Medical Report Simplifier")

uploaded_file = st.file_uploader("📄 Upload a Medical Report (PDF)", type="pdf")

if uploaded_file:
    text = extract_text_from_pdf(uploaded_file)
    st.subheader("📜 Extracted Text from Report")
    st.write(text[:1000] + " ...")  # Show first 1000 characters

    summary = summarize_text(text)
    st.subheader("📝 Simplified Medical Summary")
    st.write(summary)

    st.subheader("💬 Ask a Question About Your Report")
    user_question = st.text_input("🔍 Enter your question:")
    if user_question:
        answer = answer_question(text, user_question)
        st.write("🧑‍⚕️ **AI Answer:**", answer)