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

# Page config
st.set_page_config(page_title="🖌️ MNIST Digit Recognizer", layout="centered")

# Custom CSS styling
st.markdown("""
    <style>
        .main {
            background: linear-gradient(to right, #f7f7f7, #e3f2fd);
            font-family: 'Segoe UI', sans-serif;
        }
        .title {
            text-align: center;
            font-size: 36px;
            color: #0d47a1;
            font-weight: bold;
            margin-bottom: 10px;
        }
        .subtitle {
            text-align: center;
            font-size: 18px;
            color: #424242;
            margin-bottom: 40px;
        }
        .prediction-box {
            background-color: #e3f2fd;
            padding: 20px;
            border-radius: 15px;
            text-align: center;
            margin-top: 20px;
            box-shadow: 0 4px 6px rgba(0, 0, 0, 0.1);
        }
        .prediction-text {
            font-size: 28px;
            color: #0d47a1;
            font-weight: bold;
        }
    </style>
""", unsafe_allow_html=True)

# Title
st.markdown('<div class="title">🖌️ Handwritten Digit Recognizer</div>', unsafe_allow_html=True)
st.markdown('<div class="subtitle">Draw any digit (0–9) and let the AI predict it!</div>', unsafe_allow_html=True)

# Sidebar settings
with st.sidebar:
    st.header("🎨 Drawing Settings")
    drawing_mode = st.selectbox("Drawing 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")
    bg_image = st.file_uploader("Background Image", type=["png", "jpg"])
    realtime_update = st.checkbox("Realtime Update", True)

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

model = load_mnist_model()

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

# Prediction logic
if canvas_result.image_data is not None:
    #st.image(canvas_result.image_data, caption="🖼️ Your Drawing", use_column_width=True)
    
    # Preprocess
    img = cv2.cvtColor(canvas_result.image_data.astype("uint8"), cv2.COLOR_RGBA2GRAY)
    img = 255 - img
    img_resized = cv2.resize(img, (28, 28))
    img_normalized = img_resized / 255.0
    img_reshaped = img_normalized.reshape((1, 28, 28))

    prediction = model.predict(img_reshaped)

    st.markdown(
        f"""
        <div class="prediction-box">
            <div class="prediction-text">🔢 Predicted Digit: {np.argmax(prediction)}</div>
        </div>
        """, unsafe_allow_html=True
    )