import gradio as gr
from transformers import pipeline
import PyPDF2
import torch

# Function to extract text from PDF
def extract_text_from_pdf(pdf_file):
    text = ""
    try:
        pdf_reader = PyPDF2.PdfReader(pdf_file)
        for page in pdf_reader.pages:
            text += page.extract_text()
    except Exception as e:
        return f"Error parsing PDF: {str(e)}"
    return text

# Load Summarization Model
summarizer = pipeline("summarization", model="facebook/bart-large-cnn", device=0 if torch.cuda.is_available() else -1)

# Summarization Function
def summarize_text(text):
    try:
        summaries = summarizer(text[:1024], max_length=150, min_length=40, do_sample=False)
        return summaries[0]['summary_text']
    except Exception as e:
        return f"Error during summarization: {str(e)}"

# Load Question-Answering Model
qa_pipeline = pipeline("question-answering", model="deepset/roberta-base-squad2", device=0 if torch.cuda.is_available() else -1)

# Key Findings Extraction Function
def extract_key_findings(text, query="What are the main findings?"):
    try:
        result = qa_pipeline({'question': query, 'context': text[:1024]})
        return result['answer']
    except Exception as e:
        return f"Error extracting key findings: {str(e)}"

# Main Function to Process PDF
def process_paper(pdf_file):
    extracted_text = extract_text_from_pdf(pdf_file.name)
    if "Error" in extracted_text:
        return extracted_text, None, None
    
    summary = summarize_text(extracted_text)
    findings = extract_key_findings(extracted_text)
    
    return extracted_text, summary, findings

# Gradio Interface
with gr.Blocks() as app:
    gr.Markdown("## AI-Assisted Literature Review Tool")
    
    with gr.Row():
        pdf_input = gr.File(label="Upload Research Paper (PDF)")
        process_button = gr.Button("Process")

    with gr.Row():
        extracted_text_output = gr.Textbox(label="Extracted Text", lines=10, interactive=False)
        summary_output = gr.Textbox(label="Summary", lines=5, interactive=False)
        findings_output = gr.Textbox(label="Key Findings", lines=5, interactive=False)
    
    process_button.click(process_paper, inputs=pdf_input, outputs=[extracted_text_output, summary_output, findings_output])

app.launch()