import gradio as gr
import torch
from PIL import Image
from transformers import ViTImageProcessor, ViTForImageClassification
import spaces

REPO_NAME = "ngohjuniormbah/plasmovision-malaria-ai"

print("Loading PlasmoVision AI Core...")
processor = ViTImageProcessor.from_pretrained(REPO_NAME)
model = ViTForImageClassification.from_pretrained(REPO_NAME)

@spaces.GPU
def predict_malaria(image):
    if image is None:
        return None, "System Error: Please upload a valid blood smear micrograph."
    
    image = image.convert("RGB")
    inputs = processor(images=image, return_tensors="pt")
    
    with torch.no_grad():
        outputs = model(**inputs)
    
    logits = outputs.logits
    probabilities = torch.nn.functional.softmax(logits, dim=-1)[0]
    
    labels = model.config.id2label
    confidences = {labels[i]: float(probabilities[i]) for i in range(len(labels))}
    
    top_pred_idx = logits.argmax(-1).item()
    top_label = labels[top_pred_idx]
    top_conf = float(probabilities[top_pred_idx]) * 100
    
    if top_label.lower() == "parasitized":
        summary = f"DIAGNOSTIC STATUS: POSITIVE (Parasitized)\nConfidence Score: {top_conf:.2f}%\nClinical Protocol: High parasitemia detected. Immediate laboratory verification and medical treatment recommended."
    else:
        summary = f"DIAGNOSTIC STATUS: NEGATIVE (Uninfected)\nConfidence Score: {top_conf:.2f}%\nClinical Protocol: No Plasmodium parasites detected in the provided micrographic sample."
        
    return confidences, summary

# Custom Modern CSS matching the BeeBot Dashboard Aesthetic
custom_css = """
/* Overall Canvas */
.gradio-container {
    background-color: #f7f9fc !important;
    font-family: -apple-system, BlinkMacSystemFont, "Segoe UI", Roboto, Helvetica, Arial, sans-serif !important;
}

/* Rounded Cards with soft drop-shadows */
.block, .panel, .form {
    background: #ffffff !important;
    border-radius: 20px !important;
    border: 1px solid #edf2f7 !important;
    box-shadow: 0 4px 20px rgba(0, 0, 0, 0.02) !important;
}

/* Custom Blue Primary Buttons */
button.primary {
    background: linear-gradient(135deg, #2563eb 0%, #1d4ed8 100%) !important;
    color: #ffffff !important;
    border-radius: 14px !important;
    border: none !important;
    padding: 14px 24px !important;
    font-weight: 600 !important;
    font-size: 15px !important;
    box-shadow: 0 4px 14px rgba(37, 99, 235, 0.25) !important;
    transition: all 0.2s ease-in-out !important;
}

button.primary:hover {
    transform: translateY(-2px) !important;
    box-shadow: 0 6px 20px rgba(37, 99, 235, 0.35) !important;
}

/* Clean Labels and Inputs */
.gr-box {
    border-radius: 14px !important;
    border: 1px solid #e2e8f0 !important;
}

footer { visibility: hidden !important; }
"""

# Custom HTML for Left Sidebar Logo & Main Dashboard Header
sidebar_header = """
<div style="padding: 10px 5px; margin-bottom: 10px;">
    <div style="font-size: 26px; font-weight: 800; tracking: -0.5px;">
        <span style="color: #2563eb;">Plasmo</span><span style="color: #ef4444;">Vision</span>
    </div>
    <div style="font-size: 12px; color: #94a3b8; font-weight: 500; margin-top: 2px;">
        Clinical AI Suite v1.0
    </div>
</div>
"""

dashboard_header = """
<div style="padding: 10px 0 20px 0;">
    <h2 style="font-size: 22px; font-weight: 700; color: #0f172a; margin: 0;">Diagnostic Results Dashboard</h2>
    <p style="font-size: 14px; color: #64748b; margin-top: 4px;">Automated Microscopy Analysis System</p>
</div>
"""

theme = gr.themes.Soft(
    primary_hue="blue",
    neutral_hue="slate"
)

with gr.Blocks(css=custom_css, title="PlasmoVision | Clinical AI Suite") as demo:
    
    with gr.Row():
        # Left Panel (Sidebar Style - Inputs & Controls)
        with gr.Column(scale=1):
            gr.HTML(sidebar_header)
            
            with gr.Group():
                image_input = gr.Image(type="pil", label="Upload Micrograph")
                submit_btn = gr.Button("Run Diagnostics", variant="primary")
                
            gr.Markdown(
                """
                ---
                <div style="font-size: 12px; color: #94a3b8; text-align: center;">
                Powered by Vision Transformer (ViT)<br>
                Engineered for Cameroonian Clinical Workflows
                </div>
                """
            )
        
        # Right Panel (Dashboard Results)
        with gr.Column(scale=2):
            gr.HTML(dashboard_header)
            
            with gr.Group():
                confidence_output = gr.Label(num_top_classes=2, label="Probability Distribution")
                
            with gr.Group():
                diagnosis_output = gr.Textbox(label="Clinical Summary & Protocols", lines=4)

    # Event binding
    submit_btn.click(
        fn=predict_malaria,
        inputs=image_input,
        outputs=[confidence_output, diagnosis_output]
    )

demo.launch(theme=theme)