import os
import json
import warnings

import gradio as gr
import gdown
import numpy as np
import tensorflow as tf
import spaces

from tensorflow.keras.applications.mobilenet_v2 import preprocess_input


# ============================================================
# CONFIGURATION
# ============================================================

MODEL_FILE_ID = "10axUT2jw9rAEo8FI4i9_ZGTv83UWqMeF"
CLASS_NAMES_FILE_ID = "1ztvyL4h0sQlLeyf97E9RoZZ60yzrpmRi"
MODEL_PATH = "cataract_model.keras"
CLASS_NAMES_PATH = "class_names.json"

IMG_HEIGHT = 180
IMG_WIDTH = 180


# ============================================================
# ZEROGPU STARTUP PLACEHOLDER
# ============================================================
# ZeroGPU refuses to start a Space unless it finds at least one
# @spaces.GPU function. TensorFlow does not work well inside the
# separate GPU worker process, so this dummy function only exists
# to satisfy the startup check. The real prediction function below
# is deliberately NOT decorated and runs in the main process.

@spaces.GPU
def _zerogpu_placeholder():
    return None


# ============================================================
# GOOGLE DRIVE DOWNLOAD LOGIC
# ============================================================

def download_if_missing(file_id, dest_path, label):
    """Download a Google Drive file if it is not already present."""

    if os.path.exists(dest_path) and os.path.getsize(dest_path) > 0:
        print(f"{label} already exists. Skipping download.")
        return

    if not file_id or file_id.startswith("YOUR_"):
        raise ValueError(f"Please configure the Google Drive ID for {label}.")

    print(f"Downloading {label} from Google Drive...")

    url = f"https://drive.google.com/uc?id={file_id}"
    downloaded_path = gdown.download(url=url, output=dest_path, quiet=False)

    if (
        not downloaded_path
        or not os.path.exists(dest_path)
        or os.path.getsize(dest_path) == 0
    ):
        raise RuntimeError(
            f"Failed to download {label}. Check that the Google Drive "
            "file is shared as 'Anyone with the link'."
        )

    print(f"{label} downloaded successfully.")


download_if_missing(MODEL_FILE_ID, MODEL_PATH, "TensorFlow model")
download_if_missing(CLASS_NAMES_FILE_ID, CLASS_NAMES_PATH, "class names")


# ============================================================
# LOAD CLASS NAMES
# ============================================================

with open(CLASS_NAMES_PATH, "r", encoding="utf-8") as file:
    class_names = json.load(file)

if isinstance(class_names, dict):
    categories = [
        class_names[key]
        for key in sorted(class_names, key=lambda item: int(item))
    ]
else:
    categories = list(class_names)

if len(categories) != 2:
    raise ValueError(f"Expected two class names, but received: {categories}")

print("Class index order:", categories)


# ============================================================
# LOAD TENSORFLOW MODEL
# ============================================================

with warnings.catch_warnings():
    warnings.simplefilter("ignore")
    model = tf.keras.models.load_model(MODEL_PATH, compile=False)

print("TensorFlow version:", tf.__version__)
print("Model loaded successfully.")
print("Model input shape:", model.input_shape)


# ============================================================
# IMAGE PREPROCESSING
# ============================================================

def preprocess_image(img):
    """Prepare an uploaded PIL image for MobileNetV2."""

    img = img.convert("RGB")
    img = img.resize((IMG_WIDTH, IMG_HEIGHT))

    image_array = np.asarray(img, dtype=np.float32)
    image_array = np.expand_dims(image_array, axis=0)

    # Must match the preprocessing used during model training.
    # If your model already contains a Rescaling / preprocess_input
    # layer, remove this line to avoid applying it twice.
    return preprocess_input(image_array)


# ============================================================
# PREDICTION (runs in the main process, no @spaces.GPU)
# ============================================================

def classify_image(img):
    """Return class probabilities and a human-readable result."""

    if img is None:
        return {}, "### Upload an eye image to begin.", "—"

    image_array = preprocess_image(img)

    # Calling the model directly is faster and more reliable than
    # model.predict() for a single image.
    prediction = np.asarray(
        model(image_array, training=False)
    ).reshape(-1)

    # Supports either a single sigmoid output or a two-class output.
    if len(prediction) == 1:
        probability_class_1 = float(np.clip(prediction[0], 0.0, 1.0))
        probabilities = {
            categories[0]: 1.0 - probability_class_1,
            categories[1]: probability_class_1,
        }

    elif len(prediction) == 2:
        values = np.clip(prediction.astype(float), 0.0, None)
        total = float(values.sum())

        if total == 0:
            return {}, "### Unable to calculate a prediction.", "—"

        values = values / total
        probabilities = {categories[i]: float(values[i]) for i in range(2)}

    else:
        raise ValueError(f"Unexpected model output shape: {prediction.shape}")

    predicted_class = max(probabilities, key=probabilities.get)
    confidence = probabilities[predicted_class] * 100

    if predicted_class.strip().lower() == "cataract":
        icon = "🔎"
        description = "The model detected features associated with cataracts."
    else:
        icon = "✓"
        description = "The model classified this image as normal."

    result_markdown = f"""
## {icon} {predicted_class}

{description}

**Model confidence: {confidence:.1f}%**

*This is an AI-generated classification, not a medical diagnosis.*
"""

    return probabilities, result_markdown, f"{confidence:.1f}%"


# ============================================================
# GRADIO INTERFACE
# ============================================================

custom_css = """
.gradio-container {
    max-width: 1120px !important;
    margin: auto !important;
}

#hero {
    text-align: center;
    padding: 32px 20px 26px;
    border-radius: 22px;
    background: linear-gradient(135deg, #102b46, #145c67);
    color: white;
    margin-bottom: 18px;
}

#hero h1 {
    color: white !important;
    font-size: 2.4rem !important;
    margin-bottom: 10px;
}

#hero p {
    color: #d5eeee !important;
    font-size: 1.05rem;
}

.panel {
    border: 1px solid var(--border-color-primary);
    border-radius: 18px;
    padding: 18px;
}

#predict-button {
    background: linear-gradient(90deg, #0d9488, #0f766e);
    color: white;
    border: none;
    font-weight: 700;
    border-radius: 12px;
}

#predict-button:hover {
    filter: brightness(1.08);
}

#disclaimer {
    border-left: 4px solid #d97706;
    padding: 12px 16px;
    border-radius: 8px;
    background: var(--block-background-fill);
}
"""

theme = gr.themes.Soft(
    primary_hue="teal",
    secondary_hue="cyan",
    neutral_hue="slate",
)

with gr.Blocks(title="CataractScan AI") as demo:

    gr.HTML("""
    <div id="hero">
        <h1>👁️ CataractScan AI</h1>
        <p>
            AI-powered eye image classification using TensorFlow
            and MobileNetV2
        </p>
    </div>
    """)

    gr.Markdown(
        """
        Upload an eye image to see how the model classifies it.
        Your result includes the predicted class and the model's
        estimated probabilities.
        """
    )

    with gr.Row(equal_height=True):

        with gr.Column(scale=1, elem_classes=["panel"]):
            gr.Markdown("### 📤 Upload an eye image")

            image_input = gr.Image(
                type="pil",
                label="Eye image",
                height=340,
            )

            with gr.Row():
                predict_button = gr.Button(
                    "Analyze Image",
                    variant="primary",
                    elem_id="predict-button",
                    scale=3,
                )
                clear_button = gr.Button("Clear", scale=1)

            gr.Markdown(
                """
                **Tips:** Use a clear, well-lit eye image.
                The model may perform poorly on images that differ
                from its training data.
                """
            )

        with gr.Column(scale=1, elem_classes=["panel"]):
            gr.Markdown("### 🧠 Classification results")

            prediction_output = gr.Label(
                label="Class probabilities",
                num_top_classes=2,
            )

            result_output = gr.Markdown("### Your result will appear here.")

            confidence_output = gr.Textbox(
                label="Confidence of predicted class",
                value="—",
                interactive=False,
            )

    gr.Markdown("### 🖼️ Try an example")

    example_images = [
        [name] for name in ["cataract2.jpg","cataract.jpg", "normal.png", "mine.jpg"] if os.path.exists(name)
    ]

    if example_images:
        gr.Examples(
            examples=example_images,
            inputs=image_input,
            label="Sample eye images",
        )

    gr.Markdown(
        """
        <div id="disclaimer">

        **⚠️ Medical disclaimer**

        CataractScan AI is an educational AI project, not a
        diagnostic medical device. Its output must not be used
        to diagnose or rule out cataracts. Image quality,
        training data, and model limitations can affect results.
        Consult a qualified eye-care professional for medical
        assessment.

        </div>
        """
    )

    outputs = [prediction_output, result_output, confidence_output]

    # Only the button triggers prediction, so each image is analysed once
    # (the old .change() handler made every upload run twice).
    predict_button.click(
        fn=classify_image,
        inputs=image_input,
        outputs=outputs,
    )

    def clear_all():
        return None, {}, "### Your result will appear here.", "—"

    clear_button.click(
        fn=clear_all,
        inputs=None,
        outputs=[image_input] + outputs,
    )

if __name__ == "__main__":
    # In Gradio 6, theme and css are passed to launch(), not Blocks().
    demo.launch(theme=theme, css=custom_css)