import gradio as gr
from PIL import Image
from app.model import predict, gradcam, CLASS_NAMES

def predict_fn(img):
    label, confidence, probs = predict(img)
    probs_sorted = {k: float(v) for k, v in sorted(probs.items(), key=lambda x: x[1], reverse=True)}
    return {
        "Predicted label": label,
        "Confidence": round(confidence, 3),
        "Class probabilities": probs_sorted
    }

def gradcam_fn(img, interpolant):
    heatmap = gradcam(img, interpolant=float(interpolant))
    return Image.fromarray(heatmap)

with gr.Blocks(title="Brain Tumor MRI Classifier (InceptionV3 + Grad-CAM)") as demo:
    gr.Markdown("# Brain Tumor MRI Classifier (InceptionV3 + Grad-CAM)")
    gr.Markdown("Upload an MRI image to classify and visualize Grad-CAM explanation.")

    with gr.Row():
        with gr.Column():
            input_img = gr.Image(type="pil", label="Upload MRI Image")
            interpolant_slider = gr.Slider(0, 1, value=0.5, label="Grad-CAM Intensity (interpolant)")
            submit_btn = gr.Button("Run Prediction + Grad-CAM")

        with gr.Column():
            output_json = gr.JSON(label="Prediction Results")
            output_cam = gr.Image(label="Grad-CAM Overlay")

    submit_btn.click(
        fn=lambda img, interp: (predict_fn(img), gradcam_fn(img, interp)),
        inputs=[input_img, interpolant_slider],
        outputs=[output_json, output_cam]
    )

demo.launch()