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

# Page setup
st.set_page_config(page_title="Digit Recognizer", layout="centered")
st.markdown("<h1 style='text-align: center;'>🧠 Handwritten Digit Recognizer</h1>", unsafe_allow_html=True)
st.markdown("---")

# Sidebar settings
st.sidebar.header("🛠 Drawing Settings")
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 hex: ", "#000000")  # black
bg_color = st.sidebar.color_picker("Background color hex: ", "#FFFFFF")  # white
bg_image = st.sidebar.file_uploader("Background image:", type=["png", "jpg"])
realtime_update = st.sidebar.checkbox("Update in realtime", True)

# Load model
@st.cache_resource
def load_mnist_model():
    return load_model("mnist_model.keras")

model = load_mnist_model()

# Layout columns
col1, col2 = st.columns(2)

with col1:
    st.subheader("🖌️ Draw a Digit")
    canvas_result = st_canvas(
        fill_color="rgba(255, 165, 0, 0.3)",
        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"
    )

with col2:
    if canvas_result.image_data is not None:
        st.image(canvas_result.image_data, caption="Your Drawing")

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

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

            # st.success(f"✅ Predicted Digit: **{predicted_digit}**")

            st.markdown("<br>", unsafe_allow_html=True)
            st.markdown(
                        f"<h2 style='text-align: center; color: green;'>✅ Predicted Digit: {predicted_digit}</h2>",
            unsafe_allow_html=True
                    )
            st.markdown("<br>", unsafe_allow_html=True)