import gradio as gr
import tensorflow as tf
import numpy as np
from PIL import Image
import os # Import os module to check for file existence

# Load model
# Ensure 'Skema2.h5' is in the same directory as this script, or provide the full path.
model_path = 'best_swin_model.h5'
model = None # Initialize model as None

if not os.path.exists(model_path):
    print(f"Error: Model file '{model_path}' not found. Please ensure it's in the same directory.")
    print("Pastikan file model 'best_swin_model.h5' berada di direktori yang sama dengan script ini.")
else:
    try:
        model = tf.keras.models.load_model(model_path)
        print("Model 'best_swin_model.h5' berhasil dimuat.")
    except Exception as e:
        print(f"Error loading model from '{model_path}': {e}")
        print(f"Terjadi kesalahan saat memuat model dari '{model_path}': {e}")
        # Fallback or error handling if model cannot be loaded
        # For a production app, you might want to display an error message in the UI
        # or use a placeholder model.

# ***PERBAIKAN DI SINI: Ubah ukuran w dan h agar sesuai dengan input model Anda***
w, h = 150, 150 # Model mengharapkan input 150x150, 3 saluran warna

# Label map for the classification results
class_names = {
    0: 'glioma_tumor',
    1: 'no_tumor',
    2: 'meningioma_tumor',
    3: 'pituitary_tumor'
}

# Function to predict the tumor type from an uploaded image
def predict(image):
    # Check if the model was loaded successfully
    if model is None:
        return "Error: Model tidak dimuat. Pastikan file model (Skema2.h5) ada dan tidak rusak." \
               " (Error: Model not loaded. Please ensure the model file (Skema2.h5) exists and is not corrupted.)"

    # Handle cases where no image is provided
    if image is None:
        return "Tidak ada gambar yang diunggah. Silakan unggah gambar untuk prediksi." \
               " (No image uploaded. Please upload an image for prediction.)"

    try:
        # Convert image to RGB and resize it to the model's expected input dimensions
        img = image.convert("RGB").resize((w, h))
        # Convert the PIL image to a NumPy array and normalize pixel values to 0-1
        img_array = np.array(img) / 255.0
        # Add an extra dimension to the array to match the model's batch input shape (e.g., (1, 150, 150, 3))
        img_array = np.expand_dims(img_array, axis=0)

        # Make a prediction using the loaded model
        pred = model.predict(img_array)
        # Get the index of the highest probability prediction
        predicted_class_index = np.argmax(pred)
        # Return the corresponding class name from the label map
        return class_names[predicted_class_index]
    except Exception as e:
        # Catch any errors during prediction and return an error message
        return f"Terjadi kesalahan saat prediksi: {e}. Pastikan gambar valid dan sesuai." \
               f" (An error occurred during prediction: {e}. Please ensure the image is valid and appropriate.)"

# Gradio application layout
with gr.Blocks(
    # Custom CSS for styling the Gradio components
    css="""
    .gr-button {
        border-radius: 1rem !important; /* Rounded corners for buttons */
        background-color: #4CAF50; /* Green background */
        color: white; /* White text */
        padding: 0.75rem 1.5rem; /* Padding */
        font-size: 1.1rem; /* Larger font size */
        border: none; /* No border */
        cursor: pointer; /* Pointer cursor on hover */
        transition: background-color 0.3s ease; /* Smooth transition */
    }
    .gr-button:hover {
        background-color: #45a049; /* Darker green on hover */
    }
    .gr-image-preview {
        border-radius: 0.5rem; /* Slightly rounded corners for image preview */
        box-shadow: 0 4px 6px rgba(0,0,0,0.1); /* Subtle shadow */
        border: 1px solid #ddd; /* Light border */
    }
    .gradio-container {
        font-family: 'Inter', sans-serif; /* Use Inter font */
        max-width: 800px; /* Max width for the container */
        margin: auto; /* Center the container */
        padding: 2rem; /* Padding around the content */
        background-color: #f9f9f9; /* Light background */
        border-radius: 1rem; /* Rounded container */
        box-shadow: 0 8px 16px rgba(0,0,0,0.15); /* More prominent shadow */
    }
    h2 {
        color: #2c3e50; /* Darker heading color */
        text-align: center; /* Center align heading */
        margin-bottom: 1.5rem; /* Space below heading */
    }
    .gr-markdown p {
        text-align: center; /* Center align instructions */
        color: #555; /* Slightly darker text */
        margin-bottom: 2rem; /* Space below instructions */
    }
    .gr-textbox {
        border-radius: 0.5rem; /* Rounded text box */
        border: 1px solid #ccc; /* Light border */
        padding: 0.75rem; /* Padding inside text box */
        font-size: 1rem; /* Font size */
    }
    """
) as demo:
    # Markdown for the application title
    gr.Markdown("## 🧠 Deteksi Tumor Otak")
    # Markdown for user instructions
    gr.Markdown("Unggah gambar Tumor Otak, lalu klik tombol di bawah untuk mendapatkan prediksinya.")

    # Column layout for input and output components
    with gr.Column():
        # Image input component, allowing PIL Image type
        image = gr.Image(label="Unggah Gambar", type="pil")
        # Textbox to display the prediction result, non-interactive
        result = gr.Textbox(label="Hasil Prediksi", max_lines=1, interactive=False)
        # Button to trigger the prediction
        predict_btn = gr.Button("Kirim & Prediksi")

    # Define the action when the predict_btn is clicked:
    # Call the 'predict' function with the 'image' input and display the result in the 'result' textbox.
    predict_btn.click(fn=predict, inputs=image, outputs=result)

# Launch the Gradio application
# The 'share=True' option creates a public link for sharing (useful for demos)
# The 'debug=True' option provides more detailed error messages in the console
demo.launch()