from fastai.vision.all import *
import gradio as gr

from pathlib import Path
import pandas as pd

def get_x(row):
    # All files are assumed to be '.jpg', so we directly return the path with '.jpg' extension
    return path_image_combined / f"{row['file_name']}.jpg"

def get_y(row):
    return row['Buried ODD']

# Load the models into a dictionary
models = {
    'Ultrasound': load_learner('ODDUltrasound.pkl'),
    'OCT': load_learner('ODDOCT.pkl'),
    'Fundus': load_learner('ODDfundus.pkl'),
    'Fluorescence': load_learner('ODDfluorescence.pkl')
}

modality_keys = ['Ultrasound', 'OCT', 'Fundus', 'Fluorescence']

def classify_images(img_ultrasound=None, img_oct=None, img_fundus=None, img_fluorescence=None):
    imgs = [img_ultrasound, img_oct, img_fundus, img_fluorescence]
    predictions = []
    detailed_predictions = []  # To store detailed predictions for each modality
    
    if not any(imgs):  # Check if no images were provided
        return "Please upload at least one image for prediction."
    
    # Convert images to PILImage and predict with each model, if provided
    for img, key in zip(imgs, modality_keys):
        if img is not None:  # Only proceed if an image was uploaded for this modality
            pil_img = PILImage.create(img)
            pred, _, probs = models[key].predict(pil_img)
            predictions.append(pred)
            prob_pred = probs.max()  # Get the highest probability score
            detailed_predictions.append(f"{key}: {pred} ({prob_pred:.2f}%)")
        else:
            detailed_predictions.append(f"{key}: No image provided")
    
    # Handle final decision logic here as needed
    
    return "\n".join(detailed_predictions)


# Define the Gradio interface inputs and outputs
inputs = [gr.Image(label=f"{modality} Image", optional=True) for modality in modality_keys]
output = gr.Text(label="Predictions")

intf = gr.Interface(fn=classify_images, inputs=inputs, outputs=output,
                    title="ODD Detection from Multiple Imaging Modalities",
                    description="Upload images for each modality (as available) and receive individual predictions with percentages and a binary prediction for Optic Disk Drusen presence. At least one image is required.")
intf.launch(share=True)