import gradio as gr
from PIL import Image
import torch
from torchvision import models, transforms
import PyPDF2  # For reading patient reports (PDFs)
import io

# Load the pre-trained model (for example, ResNet18)
model = models.resnet18(pretrained=True)
model.eval()

# Define image preprocessing function
def preprocess_image(image):
    transform = transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.ToTensor(),
    ])
    return transform(image).unsqueeze(0)

# Define a simple prediction function for X-ray images
def predict_xray(image):
    image_tensor = preprocess_image(image)
    with torch.no_grad():
        outputs = model(image_tensor)
    probs = torch.nn.functional.softmax(outputs[0], dim=0)
    
    # Example output - replace with actual classes based on your model
    conditions = ["Normal", "Pneumonia", "Cancer", "TB", "Other"]
    results = {conditions[i]: float(probs[i]) for i in range(len(conditions))}
    summary = f"Summary: Based on the X-ray, the patient is diagnosed with: {max(results, key=results.get)}"
    return summary, results

# Define a function to read and analyze patient reports (PDFs)
def analyze_report(file):
    text = ""
    if file.name.endswith(".pdf"):
        pdf_reader = PyPDF2.PdfReader(file)
        for page in pdf_reader.pages:
            text += page.extract_text()
    # You can process the extracted text and provide insights
    # For now, let's assume the text contains diagnosis
    report_summary = f"Patient Report: {text[:300]}..."  # First 300 characters of report as preview
    return report_summary

# Gradio Interface
def create_interface():
    with gr.Blocks() as demo:
        # Custom CSS for UI
        custom_css = """
        .gradio-container {
            background-color: #f4f6f9;
            border-radius: 15px;
            box-shadow: 0 4px 15px rgba(0, 0, 0, 0.1);
            padding: 30px;
            font-family: 'Segoe UI', sans-serif;
        }
        .title {
            font-size: 30px;
            text-align: center;
            color: #4C6A92;
        }
        .gradio-button {
            background-color: #3B82F6;
            color: white;
            border-radius: 10px;
            padding: 15px;
        }
        .result-box {
            background-color: #ffffff;
            border-radius: 10px;
            padding: 20px;
            box-shadow: 0 4px 8px rgba(0, 0, 0, 0.1);
            margin-top: 20px;
        }
        """
        
        # Title section
        gr.Markdown("<h1 class='title'>RadiologyScan AI</h1>")
        
        # Upload X-ray image section
        with gr.Row():
            xray_input = gr.Image(label="Upload Chest X-ray", type="pil")
            report_input = gr.File(label="Upload Patient Report (PDF)", file_count="single")
        
        # Buttons for analysis
        with gr.Row():
            predict_button = gr.Button("Analyze X-ray", elem_classes="gradio-button")
            report_button = gr.Button("Analyze Report", elem_classes="gradio-button")
        
        # Results section for the X-ray image
        xray_output = gr.Textbox(label="X-ray Summary", interactive=False, elem_classes="result-box")
        xray_result = gr.JSON(label="X-ray Results", interactive=False, elem_classes="result-box")
        
        # Results section for the patient report
        report_output = gr.Textbox(label="Report Summary", interactive=False, elem_classes="result-box")
        
        # Event handlers for buttons
        predict_button.click(predict_xray, inputs=xray_input, outputs=[xray_output, xray_result])
        report_button.click(analyze_report, inputs=report_input, outputs=report_output)
    
    return demo

# Launch the Gradio interface
demo = create_interface()
demo.launch(share=True)