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

# Load the trained ANN model
model = tf.keras.models.load_model("mnist_ann.keras")


def preprocess(image):
    """Convert sketchpad/uploaded image to MNIST format:
    white digit on black, cropped to bounding box, centered in 28x28."""
    if image is None:
        return None

    # Gradio Sketchpad returns a dict with 'composite' key
    if isinstance(image, dict):
        image = image.get("composite")
        if image is None:
            return None

    img = Image.fromarray(np.array(image).astype("uint8")).convert("L")
    arr = np.array(img).astype("float32")

    # MNIST = white digit on black background.
    # Sketchpad = black drawing on white, so invert if background is white.
    if arr.mean() > 127:
        arr = 255.0 - arr

    # --- Step 1: Crop to the bounding box of the drawing ---
    mask_rows = np.any(arr > 30, axis=1)
    mask_cols = np.any(arr > 30, axis=0)
    if not mask_rows.any() or not mask_cols.any():
        return None  # empty canvas

    r0, r1 = np.where(mask_rows)[0][[0, -1]]
    c0, c1 = np.where(mask_cols)[0][[0, -1]]
    arr = arr[r0:r1 + 1, c0:c1 + 1]

    # --- Step 2: Resize to fit in 20x20 (keeping aspect ratio) ---
    h, w = arr.shape
    scale = 20.0 / max(h, w)
    new_h = max(1, int(round(h * scale)))
    new_w = max(1, int(round(w * scale)))
    img_small = Image.fromarray(arr.astype("uint8")).resize(
        (new_w, new_h), Image.LANCZOS
    )

    # --- Step 3: Paste in the center of a 28x28 black canvas ---
    canvas = np.zeros((28, 28), dtype="float32")
    top = (28 - new_h) // 2
    left = (28 - new_w) // 2
    canvas[top:top + new_h, left:left + new_w] = np.array(img_small)

    # Normalize to [0, 1]
    canvas = canvas / 255.0
    return canvas.reshape(1, 28, 28)


@spaces.GPU(duration=10)
def predict_digit(image):
    arr = preprocess(image)
    if arr is None:
        return {"Draw a digit first!": 1.0}

    probs = model.predict(arr, verbose=0)[0]
    return {str(i): float(probs[i]) for i in range(10)}


demo = gr.Interface(
    fn=predict_digit,
    inputs=gr.Sketchpad(label="Draw a digit (0-9)"),
    outputs=gr.Label(num_top_classes=3, label="Prediction"),
    title="MNIST Digit Classifier (ANN)",
    description="Assignment 3 - Model Deployment | Draw a handwritten digit and the ANN model will predict it.",
)

if __name__ == "__main__":
    demo.launch(ssr_mode=False)