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

st.set_page_config(page_title="Digit Recognizer", layout="centered")
st.title("🖌️ Handwritten Digit Recognizer")

# Sidebar settings
with st.sidebar:
    st.header("Canvas Settings")
    tool = st.selectbox("Tool", ["freedraw", "line", "rect", "circle", "transform"])
    stroke_width = st.slider("Stroke Width", 1, 25, 10)
    stroke_color = st.color_picker("Stroke Color", "#000000")
    bg_color = st.color_picker("Background Color", "#FFFFFF")
    realtime = st.checkbox("Realtime Update", True)

@st.cache_resource
def load_digit_model():
    return load_model("digit_recognisation.keras")

model = load_digit_model()

# Create canvas
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,
    height=280,
    width=280,
    drawing_mode=tool,
    key="canvas"
)

if canvas_result.image_data is not None:
    col1, col2 = st.columns(2)

    # Show original image
    with col1:
        st.subheader("Original Drawing")
        st.image(canvas_result.image_data,  use_container_width=True)

    # Preprocess image
    img = canvas_result.image_data.astype("uint8")
    gray = cv2.cvtColor(img, cv2.COLOR_RGBA2GRAY)
    inverted = 255 - gray
    resized = cv2.resize(inverted, (28, 28))
    norm_img = resized / 255.0
    reshaped = norm_img.reshape(1, 28, 28)

    # Show preprocessed image
    with col2:
        st.subheader("Processed Image")
        st.image(resized,  use_container_width=True, clamp=True)

    # Predict
    pred = model.predict(reshaped)
    digit = np.argmax(pred)

    # Display prediction
    st.markdown("---")
    st.markdown(f"<h2 style='text-align:center;'>Predicted Digit: {digit}</h2>", unsafe_allow_html=True)