import gradio as gr
import numpy as np
import cv2
from joblib import load
import os

# --- Configuration & Model Loading ---
MODEL_PATH = "orb_bow_svm.joblib"

# --- Model Constants (should match your training environment) ---
IMG_SIZE = (200, 200)
VOCAB_SIZE = 300
orb = cv2.ORB_create(nfeatures=500)

# Default classes
DEFAULT_CLASSES = [
    "Non Demented",
    "Very mild Dementia",
    "Mild Dementia",
    "Moderate Dementia"
]

# --- FINAL CORRECTED & EXPANDED EXAMPLE FILES (MUST MATCH UPLOADED FILES IN ROOT) ---
EXAMPLE_IMAGES = [
    "mild_9.jpg",
    "moderate_7.jpg",
    "non_93.jpg",
    "verymild_986.jpg",
    "moderate_36.jpg", 
    "verymild_795.jpg", 
    "verymild_8.jpg" 
]
# -----------------------------------------------------------------------------

# Attempt to load the model components
kmeans, scaler, svm = None, None, None
classes = DEFAULT_CLASSES

try:
    print(f"Attempting to load model from: {MODEL_PATH}")
    model_data = load(MODEL_PATH)
    kmeans = model_data["kmeans"]
    scaler = model_data["scaler"]
    svm = model_data["svm"]
    classes = model_data["classes"]
    print("Model loaded successfully!")
except FileNotFoundError:
    print(f"ERROR: Model file '{MODEL_PATH}' not found.")
except Exception as e:
    print(f"ERROR: An unexpected error occurred during model loading: {e}.")

# Define real/dummy functions based on successful load
if svm is None:
    def encode(descriptors, kmeans_model): return np.zeros(VOCAB_SIZE)
    def gradio_predict(input_img):
        return "⚠️ Model not loaded. Cannot perform prediction.", {cls: 0.0 for cls in DEFAULT_CLASSES}
else:
    def encode(descriptors, kmeans_model):
        if descriptors is None or len(descriptors) == 0:
            return np.zeros(VOCAB_SIZE)
        words = kmeans_model.predict(descriptors)
        hist, _ = np.histogram(words, bins=np.arange(VOCAB_SIZE + 1))
        return hist
    
    def gradio_predict(input_img):
        # Preprocessing
        img = cv2.cvtColor(input_img, cv2.COLOR_RGB2GRAY)
        img = cv2.resize(img, IMG_SIZE)

        # Feature Extraction
        kps, des = orb.detectAndCompute(img, None)
        feat = encode(des, kmeans).reshape(1, -1)
        feat_scaled = scaler.transform(feat)

        # Prediction
        prediction_index = svm.predict(feat_scaled)[0]
        probabilities = svm.predict_proba(feat_scaled)[0]
        
        # Format output
        predicted_class = classes[prediction_index]
        confidence_score = probabilities[prediction_index] * 100
        
        output_message = f"**Diagnosis: {predicted_class}**\n" \
                         f"Confidence: {confidence_score:.2f}%"
        
        prob_dict = {cls: prob for cls, prob in zip(classes, probabilities)}
        
        return output_message, prob_dict

# --- Gradio Interface Definition ---

colorful_theme = gr.Theme.from_hub("gradio/seafoam")

with gr.Blocks(theme=colorful_theme, title="DementiaScan-Predict 🧠") as demo:
    
    # ------------------------------------------------
    # INTRODUCTORY TEXT (UPDATED)
    # ------------------------------------------------
    gr.Markdown(
        """
        # 🧠 DementiaScan-Predict: Rapid Stage Classification 🌟
        
        Welcome! This tool offers **rapid, preliminary classification of dementia stages** from MRI brain scans. Our core innovation is providing highly **efficient and accessible AI diagnostics**, perfect for deployment in resource-constrained environments.
        
        ---
        ### 🚀 How It Works:
        1.  **Upload an MRI Scan** (T1-weighted image).
        2.  **Click 'Classify Scan'** to trigger the analysis.
        3.  **Get Instant Results** for the predicted dementia stage and confidence.
        
        ---
        """
    )
    
    with gr.Row(variant="panel"):
        # Input Column
        with gr.Column(scale=1):
            gr.Markdown("## 📤 Input Image")
            image_input = gr.Image(
                type="numpy", 
                label="Upload MRI Brain Scan Image", 
                height=350,
                width=350,
                interactive=True
            )
            
            # Action Button
            submit_btn = gr.Button("✨ Classify Scan ✨", variant="primary", size="lg")
            gr.Markdown("---")
            gr.Markdown(
                """
                ### 💡 Quick Test Examples:
                Click on any image below to load and classify it instantly!
                """
            )
            
            gr.Examples(
                examples=EXAMPLE_IMAGES,
                inputs=image_input,
                outputs=[gr.Textbox(), gr.Label()],
                fn=gradio_predict,
                cache_examples=True,
            )


        # Output Column
        with gr.Column(scale=2):
            gr.Markdown("## ✅ Prediction Results")
            output_text = gr.Textbox(
                label="Predicted Dementia Stage & Confidence", 
                value="Upload an image and click 'Classify Scan' to see the results.",
                lines=3,
                show_copy_button=True,
                elem_id="prediction_output_box" 
            )
            
            output_label = gr.Label(
                label="Detailed Probability Distribution",
            )
            
            # ------------------------------------------------
            # METHODOLOGY TEXT (UPDATED)
            # ------------------------------------------------
            gr.Markdown(
                """
                ---
                ### 📚 Methodology: ORB-BoVW-SVM
                We employ a fast, classical Computer Vision pipeline for efficiency:
                * **Feature Detection:** **ORB** detects key visual points on the brain scan.
                * **Feature Encoding:** **Bag of Visual Words (BoVW)** converts these features into a compact, fixed-size histogram (vector).
                * **Classification:** The resulting vector is classified using a **Support Vector Machine (SVM)**.
                
                This approach ensures excellent **speed and low computational overhead** compared to standard deep learning models.
                """
            )
    
    # Connect the button to the prediction function
    submit_btn.click(
        fn=gradio_predict,
        inputs=[image_input],
        outputs=[output_text, output_label]
    )

# Launch the Gradio app 
demo.launch(share=True)