import gradio as gr
from transformers import pipeline
from PIL import Image
import pytesseract
from PyPDF2 import PdfReader
from reportlab.lib.pagesizes import letter
from reportlab.pdfgen import canvas
import os

# Load the medical analysis model (e.g., BioGPT or PubMedBERT)
medical_analyzer = pipeline("text-classification", model="microsoft/biogpt")

# Function to extract text from images or PDFs
def extract_text(file_path):
    if file_path.endswith(".pdf"):
        # Extract text from PDF
        reader = PdfReader(file_path)
        text = ""
        for page in reader.pages:
            text += page.extract_text()
        return text.strip()
    else:
        # Extract text from image
        return pytesseract.image_to_string(Image.open(file_path))

# Function to generate a PDF report
def create_pdf_report(analysis, output_path):
    c = canvas.Canvas(output_path, pagesize=letter)
    c.drawString(100, 750, "Blood Test Report Analysis")
    c.drawString(100, 730, "---------------------------")
    
    y_position = 700
    for line in analysis.split("\n"):
        c.drawString(100, y_position, line)
        y_position -= 20  # Move down for the next line
    
    c.save()
    return output_path

# Function to analyze blood test reports
def analyze_blood_test(file):
    # Step 1: Extract text
    extracted_text = extract_text(file)
    if not extracted_text:
        return "Could not extract text. Please upload a valid file.", None

    # Step 2: Use medical model to analyze extracted text
    analysis_results = medical_analyzer(extracted_text)
    analysis_report = "🔍 Analysis Results:\n"
    for item in analysis_results[:5]:  # Limit results for simplicity
        analysis_report += f"- {item['label']}: {item['score']:.2f}\n"

    # Step 3: Generate downloadable PDF report
    output_pdf = "analysis_report.pdf"
    create_pdf_report(f"Extracted Text:\n{extracted_text}\n\n{analysis_report}", output_pdf)
    
    return analysis_report, output_pdf

# Gradio interface setup
interface = gr.Interface(
    fn=analyze_blood_test,
    inputs=gr.File(label="Upload your blood test report (PNG, JPG, JPEG, or PDF)"),
    outputs=[
        gr.Textbox(label="Analysis Results"),
        gr.File(label="Download PDF Report")
    ],
    title="MedAI Analyzer",
    description=(
        "Upload your blood test report in image (PNG, JPG, JPEG) or PDF format. "
        "The app will extract and analyze the values, flag abnormalities, and provide health recommendations. "
        "You can also download a detailed PDF report of the analysis."
    ),
    allow_flagging="never"
)

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