import gradio as gr
import numpy as np
from app.predictor import predict_disease, compute_ai_risk
from app.gradcam import generate_gradcam


def predict(image):
    if image is None:
        return "Upload image first", None

    label, confidence, tensor, original = predict_disease(image)

    # GradCAM
    heatmap = generate_gradcam(tensor, label)

    # ensure numpy
    heatmap = np.array(heatmap)

    # safe normalization
    heatmap = heatmap.astype(np.float32)

    if heatmap.max() > 1:
        heatmap = heatmap / 255.0

    heatmap = np.clip(heatmap, 0, 1)
    heatmap = (heatmap * 255).astype(np.uint8)

    # ensure RGB format
    if len(heatmap.shape) == 2:
        heatmap = np.stack([heatmap]*3, axis=-1)

    risk = compute_ai_risk(label, confidence, heatmap)

    report = (
        "🫀 CardioGuard AI Report\n\n"
        f"Disease: {label}\n"
        f"Confidence: {confidence*100:.2f}%\n"
        f"Risk: {risk}"
    )

    return report, heatmap


demo = gr.Blocks()

with demo:
    gr.Markdown("# 🫀 CardioGuard AI")

    with gr.Row():
        img = gr.Image(type="pil")
        out2 = gr.Image()

    out1 = gr.Textbox(label="Report")

    btn = gr.Button("Analyze")

    btn.click(
        fn=predict,
        inputs=img,
        outputs=[out1, out2]
    )


demo.queue()

demo.launch(
    server_name="0.0.0.0",
    server_port=7860
)