import gradio as gr
from transformers import AutoImageProcessor, AutoModelForImageClassification
from PIL import Image, ImageDraw, ImageFont
import torch
import numpy as np

model_name = "Hemgg/brain-tumor-classification"
processor = AutoImageProcessor.from_pretrained(model_name)
model = AutoModelForImageClassification.from_pretrained(model_name)
class_names = ["🧠 Glioma", "🎯 Meningioma", "✅ No Tumor", "⚡ Pituitary"]

custom_css = """
.gradio-container {background-color: #050505 !important; color: #ffffff !important;}
.md, .md p, .md h1, .md h2, .md h3, span, label {color: #ffffff !important;}
#result-box {
    background-color: #1a1a1a !important; 
    border: 2px solid #3b82f6 !important; 
    border-radius: 12px; 
    padding: 20px;
    color: #ffffff !important;
}
.gr-button-primary {
    background: linear-gradient(135deg, #1e40af, #7e22ce) !important;
    border: none !important;
    font-weight: bold !important;
}
#result-text * { color: #ffffff !important; }
footer {display: none !important;}
"""

def get_medical_info(tumor_type):
    info = {
        "🧠 Glioma": {
            "desc": "Gliomas originate in the glial cells that support neurons. They can be fast-growing and may involve surrounding brain tissue.",
            "next": "Urgent consultation with a neuro-oncologist. An MRI with contrast or a biopsy is typically the next diagnostic step."
        },
        "🎯 Meningioma": {
            "desc": "These tumors arise from the meninges, the layers covering the brain. Most are slow-growing and benign but can cause pressure.",
            "next": "Consult a neurosurgeon to evaluate the mass effect. Treatment ranges from 'watchful waiting' to surgical resection."
        },
        "⚡ Pituitary": {
            "desc": "Pituitary adenomas occur in the master gland at the base of the brain. They often affect hormone regulation and vision.",
            "next": "An endocrine workup (blood tests) and a visual field test are recommended to assess hormonal and optic nerve impact."
        },
        "✅ No Tumor": {
            "desc": "The neural network did not detect significant signs of the three primary tumor types in this scan.",
            "next": "If symptoms (headaches, seizures, vision loss) persist, please consult a neurologist for a comprehensive evaluation."
        }
    }
    return info.get(tumor_type, {"desc": "", "next": ""})

def classify_tumor(image):
    if image is None:
        return "<p style='color:white;'>Please upload a scan.</p>", None
    
    inputs = processor(images=image, return_tensors="pt")
    with torch.no_grad():
        outputs = model(**inputs)
        probs = torch.nn.functional.softmax(outputs.logits, dim=-1)
        idx = probs.argmax(-1).item()
        conf = probs[0][idx].item()
        name = class_names[idx]
        med = get_medical_info(name)

    overlay = create_overlay(image, idx, conf)
    
    html_res = f"""
    <div id="result-box">
        <h2 style="color: #60a5fa; margin-top:0;">{name} Detected</h2>
        <p style="font-size: 1.2em;"><b>Confidence Score:</b> <span style="color: #4ade80;">{conf:.1%}</span></p>
        <div style="margin: 15px 0; padding: 10px; background: #262626; border-radius: 8px;">
            <p style="color: #93c5fd; margin-bottom: 5px;"><b>Clinical Overview:</b></p>
            <p style="color: #ffffff;">{med['desc']}</p>
        </div>
        <div style="margin: 15px 0; padding: 10px; background: #1e3a8a; border-radius: 8px;">
            <p style="color: #bfdbfe; margin-bottom: 5px;"><b>Recommended Next Steps:</b></p>
            <p style="color: #ffffff;">{med['next']}</p>
        </div>
        <p style="font-size: 0.8em; color: #94a3b8; border-top: 1px solid #444; pt-10;">
            Probabilities: {", ".join([f"{class_names[i]}: {probs[0][i]:.1%}" for i in range(4)])}
        </p>
    </div>
    """
    return html_res, overlay

def create_overlay(image, idx, conf):
    overlay = image.copy().convert("RGBA")
    draw = ImageDraw.Draw(overlay)
    if idx != 2:
        w, h = overlay.size
        r = min(w, h) // 4
        cx, cy = w // 2, h // 2
        draw.ellipse([cx-r, cy-r, cx+r, cy+r], fill=(255, 0, 0, int(150*conf)), outline=(255, 255, 0, 255), width=5)
    return overlay.convert("RGB")

with gr.Blocks(css=custom_css, theme=gr.themes.Default()) as demo:
    gr.HTML("<h1 style='text-align:center; color:white; font-size: 32px;'>NEURO-DIAGNOSTIC AI STATION</h1>")
    
    with gr.Row():
        with gr.Column(scale=1):
            img_input = gr.Image(label="Upload Patient MRI", type="pil")
            run_btn = gr.Button("PERFORM NEURAL SCAN", variant="primary")
            
        with gr.Column(scale=1):
            res_html = gr.HTML(label="Diagnostic Findings", elem_id="result-text")
            img_output = gr.Image(label="Visualization Overlay")
            
    run_btn.click(classify_tumor, img_input, [res_html, img_output])

if __name__ == "__main__":
    demo.launch()