import os
import numpy as np
import tensorflow as tf
import gradio as gr
from PIL import Image

print("=" * 50)
print("🚀 Starting Digit Detector")
print("=" * 50)

# ----------------------------
# Load Model
# ----------------------------
MODEL_PATH = "digit_detector.keras"

if not os.path.exists(MODEL_PATH):
    raise FileNotFoundError(f"Model not found: {MODEL_PATH}")

model = tf.keras.models.load_model(MODEL_PATH)

print("✅ Model loaded successfully!")

# ----------------------------
# Prediction Function
# ----------------------------
def predict_digit(image):
    if image is None:
        return {
            "error": "Please upload an image."
        }

    try:
        # Convert numpy image to PIL
        if len(image.shape) == 3:
            img = Image.fromarray(image.astype("uint8")).convert("L")
        else:
            img = Image.fromarray(image.astype("uint8"))

        # Resize to MNIST size
        img = img.resize((28, 28))

        # Convert to array
        img = np.array(img)

        # Invert colors
        img = 255 - img

        # Normalize
        img = img.astype("float32") / 255.0

        # Shape for CNN
        img = img.reshape(1, 28, 28, 1)

        # Predict
        prediction = model.predict(img, verbose=0)

        digit = int(np.argmax(prediction))
        confidence = float(np.max(prediction) * 100)

        probs = {
            str(i): round(float(prediction[0][i]) * 100, 2)
            for i in range(10)
        }

        return {
            "Predicted Digit": digit,
            "Confidence (%)": round(confidence, 2),
            "Probabilities (%)": probs
        }

    except Exception as e:
        return {
            "error": str(e)
        }

# ----------------------------
# Gradio UI
# ----------------------------
with gr.Blocks(title="Digit Detector") as demo:

    gr.Markdown("# ✍️ Handwritten Digit Detector")
    gr.Markdown(
        "Upload an image containing a handwritten digit (0-9)."
    )

    with gr.Row():
        image = gr.Image(
            type="numpy",
            label="Upload Image"
        )

    predict_btn = gr.Button("Predict")

    output = gr.JSON(label="Prediction")

    predict_btn.click(
        fn=predict_digit,
        inputs=image,
        outputs=output
    )

print("✅ Launching Gradio...")

demo.launch(
    server_name="0.0.0.0",
    server_port=7860
)