# Install necessary libraries
#!pip install gradio transformers pandas PyPDF2 pdfplumber torch torchvision timm sentencepiece

import gradio as gr
from transformers import pipeline
import pandas as pd
import PyPDF2
import pdfplumber
import torch
import timm
from PIL import Image

# Load pre-trained model for zero-shot classification
classifier = pipeline("zero-shot-classification", model="facebook/bart-large-mnli")

# Pre-trained model for X-ray analysis (example with a model from timm library)
image_model = timm.create_model('resnet50', pretrained=True)
image_model.eval()

# Initialize patient database
patients_db = []

# Function to register patients
def register_patient(name, age, gender):
    patient_id = len(patients_db) + 1
    patients_db.append({
        "ID": patient_id, 
        "Name": name, 
        "Age": age, 
        "Gender": gender, 
        "Symptoms": "", 
        "Diagnosis": "", 
        "Action Plan": "", 
        "Medications": "", 
        "Tests": ""
    })
    return f"✅ Patient {name} registered successfully. Patient ID: {patient_id}"

# Function to analyze text reports
def analyze_report(patient_id, report_text):
    candidate_labels = ["anemia", "viral infection", "liver disease", "kidney disease", "diabetes"]
    result = classifier(report_text, candidate_labels)
    diagnosis = result['labels'][0]
    action_plan = f"Based on your report, you might have {diagnosis}. Please consult a doctor for confirmation."
    
    # Store diagnosis in the database
    for patient in patients_db:
        if patient["ID"] == patient_id:
            patient["Diagnosis"] = diagnosis
            patient["Action Plan"] = action_plan
            break
    
    return f"🔍 Diagnosis: {diagnosis}. {action_plan}"

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

# Function to analyze uploaded images (X-ray/CT-scan)
def analyze_image(patient_id, img):
    image = Image.open(img).convert('RGB')
    transform = torch.nn.Sequential(
        torch.nn.Upsample(size=(224, 224)),
        torch.nn.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    )
    image_tensor = transform(torch.unsqueeze(torch.tensor(image), 0))

    # Run the image through the model (for simplicity, assuming ResNet50 output)
    output = image_model(image_tensor)
    _, predicted = torch.max(output, 1)

    # Map prediction to a label (this is just a placeholder example)
    labels = {0: "Normal", 1: "Pneumonia", 2: "Liver Disorder", 3: "COVID-19"}
    diagnosis = labels.get(predicted.item(), "Unknown")

    # Store diagnosis in the database
    for patient in patients_db:
        if patient["ID"] == patient_id:
            patient["Diagnosis"] = diagnosis
            break
    
    return f"🔍 Diagnosis from image: {diagnosis}"

# Function to display the dashboard
def show_dashboard():
    if not patients_db:
        return "No patient records available."
    return pd.DataFrame(patients_db)

# Gradio interface for patient registration
patient_interface = gr.Interface(
    fn=register_patient,
    inputs=[
        gr.Textbox(label="Patient Name", placeholder="Enter the patient's full name"),
        gr.Number(label="Age"),  
        gr.Radio(label="Gender", choices=["Male", "Female", "Other"])
    ],
    outputs="text",
    description="Register a new patient"
)

# Gradio interface for report analysis (text input)
report_interface = gr.Interface(
    fn=analyze_report,
    inputs=[
        gr.Number(label="Patient ID"),  
        gr.Textbox(label="Report Text", placeholder="Paste the text from your report here")
    ],
    outputs="text",
    description="Analyze blood, LFT, or other medical reports"
)

# Gradio interface for PDF report analysis (PDF upload)
pdf_report_interface = gr.Interface(
    fn=lambda pdf: extract_pdf_report(pdf),
    inputs=gr.File(label="Upload PDF Report"),
    outputs="text",
    description="Extract and analyze text from PDF reports"
)

# Gradio interface for X-ray/CT-scan image analysis
image_interface = gr.Interface(
    fn=analyze_image,
    inputs=[
        gr.Number(label="Patient ID"),
        gr.Image(type="filepath", label="Upload X-ray or CT-Scan Image")
    ],
    outputs="text",
    description="Analyze X-ray or CT-scan images for diagnosis"
)

# Gradio interface for the dashboard
dashboard_interface = gr.Interface(
    fn=show_dashboard,
    inputs=None,
    outputs="dataframe",
    description="View patient reports and history"
)

# Organize the layout using Blocks
with gr.Blocks() as demo:
    gr.Markdown("# Medical Report and Image Analyzer")
    with gr.TabItem("Patient Registration"):
        patient_interface.render()
    with gr.TabItem("Analyze Report (Text)"):
        report_interface.render()
    with gr.TabItem("Analyze Report (PDF)"):
        pdf_report_interface.render()
    with gr.TabItem("Analyze Image (X-ray/CT)"):
        image_interface.render()
    with gr.TabItem("Dashboard"):
        dashboard_interface.render()

demo.launch(share=True)