import streamlit as st
from streamlit_drawable_canvas import st_canvas
from keras.models import load_model
import numpy as np
import cv2

# 🎯 App title and intro
st.set_page_config(page_title="MNIST Digit Recognizer", layout="centered")
st.title("🤖 VisionAI: MNIST Digit Predictor")
st.markdown("Draw a digit (0-9) below and click **Predict** to see the result!")

# 🎨 Sidebar controls
drawing_mode = st.sidebar.selectbox("Drawing tool:", ("freedraw", "line", "rect", "circle", "transform"))
stroke_width = st.sidebar.slider("Stroke width", 1, 25, 10)
stroke_color = st.sidebar.color_picker("Stroke color", "#000000")
bg_color = st.sidebar.color_picker("Background color", "#FFFFFF")
bg_image = st.sidebar.file_uploader("Background image (optional):", type=["png", "jpg"])
realtime_update = st.sidebar.checkbox("Update in realtime", True)

# 🧠 Load model from local path
@st.cache_resource
def load_mnist_model():
    return load_model("digit_reco.keras")

model = load_mnist_model()

# 🖌️ Canvas setup
canvas_result = st_canvas(
    fill_color="rgba(255, 165, 0, 0.3)",  # Transparent fill
    stroke_width=stroke_width,
    stroke_color=stroke_color,
    background_color=bg_color,
    update_streamlit=realtime_update,
    height=280,
    width=280,
    drawing_mode=drawing_mode,
    key="canvas",
)

# 📤 Predict button
if st.button("🔮 Predict"):
    if canvas_result.image_data is not None:
        st.image(canvas_result.image_data, caption="🖼️ Your Drawing", use_container_width=True)  # Updated

        # Preprocess image
        img = cv2.cvtColor(canvas_result.image_data.astype("uint8"), cv2.COLOR_RGBA2GRAY)
        img = 255 - img  # Invert colors
        img_resized = cv2.resize(img, (28, 28))
        img_normalized = img_resized / 255.0
        img_reshaped = img_normalized.reshape((1, 28, 28))

        # Model prediction
        prediction = model.predict(img_reshaped)
        predicted_digit = np.argmax(prediction)

        # Show result
        st.markdown(f"## 🧠 Predicted Digit: **{predicted_digit}**")
    else:
        st.warning("Please draw something before predicting!")