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

# Page configuration
st.set_page_config(page_title="Digit Recognizer", layout="centered")

# Load trained model (preferably CNN-based on MNIST)
@st.cache_resource
def load_mnist_model():
    return load_model("mnist_model.keras")

model = load_mnist_model()

# Custom CSS Styling
st.markdown("""
    <style>
    .main-title {
        text-align: center;
        font-size: 40px;
        font-weight: 700;
        color: #2c3e50;
    }
    .subtitle {
        text-align: center;
        font-size: 18px;
        color: #555;
        margin-bottom: 20px;
    }
    .result-box {
        background-color: #f0f9ff;
        border: 2px solid #3498db;
        border-radius: 10px;
        padding: 15px;
        text-align: center;
    }
    .digit {
        font-size: 36px;
        color: #2c3e50;
        font-weight: bold;
    }
    </style>
""", unsafe_allow_html=True)

st.markdown('<div class="main-title">✍️ Digit Recognizer</div>', unsafe_allow_html=True)
st.markdown('<div class="subtitle">Draw any digit (0-9) below and let the model predict it</div>', unsafe_allow_html=True)

# Sidebar controls
st.sidebar.header("🛠️ Canvas Settings")
stroke_width = st.sidebar.slider("Stroke Width", 5, 25, 15)
stroke_color = st.sidebar.color_picker("Stroke Color", "#000000")
bg_color = st.sidebar.color_picker("Background Color", "#FFFFFF")
realtime = st.sidebar.checkbox("Update in Realtime", True)

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

# Preprocess drawing like MNIST
def preprocess_drawn_image(img_data):
    gray = cv2.cvtColor(img_data.astype("uint8"), cv2.COLOR_RGBA2GRAY)
    gray = 255 - gray  # Invert to white digit on black
    _, thresh = cv2.threshold(gray, 50, 255, cv2.THRESH_BINARY)

    contours, _ = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    if not contours:
        return None

    x, y, w, h = cv2.boundingRect(contours[0])
    digit = thresh[y:y+h, x:x+w]

    # Center the digit in a square image
    max_dim = max(w, h)
    square = np.zeros((max_dim, max_dim), dtype=np.uint8)
    x_offset = (max_dim - w) // 2
    y_offset = (max_dim - h) // 2
    square[y_offset:y_offset+h, x_offset:x_offset+w] = digit

    # Resize to 20x20, then embed in 28x28
    resized = cv2.resize(square, (20, 20))
    final = np.zeros((28, 28), dtype=np.uint8)
    final[4:24, 4:24] = resized

    final = final / 255.0
    return final.reshape(1, 28, 28, 1)

# Predict and display result
if canvas_result.image_data is not None:
    processed_img = preprocess_drawn_image(canvas_result.image_data)

    if processed_img is not None:
        st.image(processed_img.reshape(28, 28), caption="🧼 Preprocessed Image", clamp=True, channels="GRAY")

        prediction = model.predict(processed_img)
        pred_digit = int(np.argmax(prediction))
        confidence = float(np.max(prediction)) * 100

        st.markdown(f"""
            <div class='result-box'>
                🧠 Predicted Digit: <span class='digit'>{pred_digit}</span><br>
                📊 Confidence: <strong>{confidence:.2f}%</strong>
            </div>
        """, unsafe_allow_html=True)
    else:
        st.warning("Couldn't detect a digit. Please try drawing again.")