import os
import urllib.request
import numpy as np
import tensorflow as tf
from tensorflow.keras.applications import EfficientNetB0
from tensorflow.keras.layers import Dense, GlobalAveragePooling2D, Dropout, BatchNormalization
from tensorflow.keras.models import Model
import gradio as gr
import cv2
import spaces  # <--- ต้องมี import นี้

MODEL_PATH = "efficientnetb0_finetuned_brain_mri.keras"
MODEL_URL = "https://huggingface.co/starpreeda/BrainTumorTest/resolve/main/efficientnetb0_finetuned_brain_mri.keras"

def load_brain_mri_model():
    if not os.path.exists(MODEL_PATH):
        print("Downloading model weights from Hugging Face Repository...")
        try:
            urllib.request.urlretrieve(MODEL_URL, MODEL_PATH)
            print("Model downloaded successfully!")
        except Exception as e:
            raise RuntimeError(f"ไม่สามารถดาวน์โหลดโมเดลได้: {e}") from e

    base_model = EfficientNetB0(weights=None, include_top=False, input_shape=(224, 224, 3))
    x = base_model.output
    x = GlobalAveragePooling2D()(x)
    x = BatchNormalization()(x)
    x = Dense(256, activation='relu')(x)
    x = Dropout(0.4)(x)
    outputs = Dense(4, activation='softmax')(x)

    model = Model(inputs=base_model.input, outputs=outputs)
    model.load_weights(MODEL_PATH)
    return model

model = load_brain_mri_model()

CLASS_MAPPING = {
    'glioma': {'name': 'Glioma Tumor', 'desc': 'A type of tumor that originates in the glial cells.'},
    'meningioma': {'name': 'Meningioma Tumor', 'desc': 'A tumor arising from the meninges.'},
    'notumor': {'name': 'No Tumor Detected', 'desc': 'No clear evidence of brain tumor tissue.'},
    'pituitary': {'name': 'Pituitary Tumor', 'desc': 'An abnormal growth located in the pituitary gland.'}
}
CLASS_NAMES = ['glioma', 'meningioma', 'notumor', 'pituitary']

# ใส่ @spaces.GPU กลับเข้ามาเพื่อรองรับ ZeroGPU environment
@spaces.GPU
def predict_mri(input_img):
    if input_img is None:
        return "<h3 style='color:#d93025;'>Please upload a valid Brain MRI scan image.</h3>", {}

    if input_img.ndim == 2:
        input_img = cv2.cvtColor(input_img, cv2.COLOR_GRAY2RGB)
    elif input_img.shape[-1] == 4:
        input_img = cv2.cvtColor(input_img, cv2.COLOR_RGBA2RGB)

    img_resized = cv2.resize(input_img, (224, 224))
    img_array = img_resized.astype(np.float32)
    img_batch = np.expand_dims(img_array, axis=0)

    predictions = model.predict(img_batch, verbose=0)[0]

    confidences = {}
    for idx, class_key in enumerate(CLASS_NAMES):
        confidences[CLASS_MAPPING[class_key]['name']] = float(predictions[idx])

    top_idx = int(np.argmax(predictions))
    top_key = CLASS_NAMES[top_idx]
    top_confidence = predictions[top_idx] * 100
    info = CLASS_MAPPING[top_key]

    summary_html = f"""
    <div style="background-color: #f8f9fa; border-left: 6px solid #1a73e8; padding: 18px; border-radius: 8px;">
        <h3 style="color: #1a73e8; margin-top: 0;">Diagnostic Classification Summary</h3>
        <p style="font-size: 20px; font-weight: bold;">Predicted Class: <span style="color: #d93025;">{info['name']}</span></p>
        <p style="font-size: 16px; font-weight: bold;">Confidence Score: <span style="color: #188038;">{top_confidence:.2f}%</span></p>
        <hr style="border: 0.5px solid #dadce0;">
        <p style="font-size: 14px; color: #5f6368;"><b>Clinical Note:</b> {info['desc']}</p>
    </div>
    """
    return summary_html, confidences

demo = gr.Interface(
    fn=predict_mri,
    inputs=gr.Image(type="numpy", label="Upload Brain MRI Image"),
    outputs=[
        gr.HTML(label="Classification Result"),
        gr.Label(num_top_classes=4, label="Class Probability Distribution")
    ],
    title="🧠 Brain Tumor MRI Classification System",
    description="Upload a Brain MRI scan to analyze potential tumor types."
)

if __name__ == "__main__":
    demo.queue().launch(server_name="0.0.0.0", server_port=7860)