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

# Page configuration
st.set_page_config(page_title="DigitGlow - AI Digit Identifier", layout="centered")

# Custom CSS Styling
st.markdown(
    """
    <style>
        .stApp {
            background-color: #1e1e1e;
            font-family: 'Segoe UI', sans-serif;
        }
        .title-box {
            text-align: center;
            padding: 20px;
        }
        .main-title {
            font-size: 3em;
            font-weight: bold;
            color: #00ffcc;
            text-shadow: 1px 1px 5px #00ffff;
        }
        .sub-title {
            font-size: 1.2em;
            color: #ccc;
            font-style: italic;
        }
        .prediction-box {
            text-align: center;
            margin-top: 30px;
        }
        .prediction-text {
            font-size: 3em;
            font-weight: bold;
            color: #ff4d4d;
            text-shadow: 2px 2px 8px #ff0000;
        }
        .canvas-title {
            text-align: center;
            color: #00ccff;
            font-size: 1.2em;
            margin-bottom: 10px;
        }
        .sidebar .sidebar-content {
            background-color: #2c2c2c;
        }
    </style>
    <div class="title-box">
        <div class="main-title">DigitGlow</div>
        <div class="sub-title">AI-Powered Handwritten Digit Recognition</div>
    </div>
    <hr style="border-top: 1px solid #555;">
    """,
    unsafe_allow_html=True
)

# Sidebar - Drawing Settings
st.sidebar.header("✏️ Canvas 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:", "#FFFFFF")
bg_color = st.sidebar.color_picker("Canvas background:", "#000000")
realtime_update = st.sidebar.checkbox("Update in real-time", True)

# Load Model
@st.cache_resource
def load_mnist_model():
    return load_model("handwritten_digit_recognition.keras")

model = load_mnist_model()

# Layout columns
col1, col2 = st.columns([1, 1])

with col1:
    st.markdown('<div class="canvas-title">🖌️ Draw a digit below:</div>', unsafe_allow_html=True)
    canvas_result = st_canvas(
        fill_color="rgba(255, 255, 255, 1)",
        stroke_width=stroke_width,
        stroke_color=stroke_color,
        background_color=bg_color,
        update_streamlit=realtime_update,
        height=280,
        width=250,
        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", width=280)
        # Preprocess image
        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)

        # Display Prediction
        st.markdown(
            f'<div class="prediction-box"><div class="prediction-text">Prediction: {np.argmax(prediction)}</div></div>',
            unsafe_allow_html=True,
        )