import os
os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"

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, ImageOps

# Custom styles
st.markdown("""
    <style>
    .big-font {
        font-size:40px !important;
        font-weight: bold;
        color: #5A189A;
        text-align: center;
    }
    .result-box {
        background-color: #F0EBF8;
        border-radius: 10px;
        padding: 20px;
        text-align: center;
        font-size: 24px;
        color: #3C096C;
        font-weight: bold;
    }
    </style>
""", unsafe_allow_html=True)

# App title
st.markdown('<p class="big-font">✍️ Handwritten Digit Recognizer</p>', unsafe_allow_html=True)
st.markdown("""
### 🖼️ Draw or Upload a digit and let the model 🤖 identify it with confidence!
""")

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

model = load_mnist_model()

# Preprocessing
def preprocess(img):
    img = ImageOps.grayscale(img)
    img = img.resize((200, 200))
    img = np.array(img)
    if np.mean(img) > 127:
        img = 255 - img
    _, img = cv2.threshold(img, 100, 255, cv2.THRESH_BINARY)
    coords = cv2.findNonZero(img)
    if coords is not None:
        x, y, w, h = cv2.boundingRect(coords)
        digit = img[y:y+h, x:x+w]
    else:
        return np.zeros((1, 28, 28), dtype="float32")
    digit = cv2.resize(digit, (20, 20), interpolation=cv2.INTER_AREA)
    digit = np.pad(digit, ((4, 4), (4, 4)), mode="constant", constant_values=0)
    digit = digit.astype("float32") / 255.0
    digit = digit.reshape(1, 28, 28)
    return digit

# Sidebar settings
st.sidebar.title("🛠️ Drawing Settings")
mode = st.sidebar.selectbox("✏️ Drawing Tool", ("freedraw", "line"))
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")

# Tabs for draw vs upload
tab1, tab2 = st.tabs(["🖌️ Draw Digit", "📤 Upload Image"])

input_img = None

# Tab 1: Draw
with tab1:
    canvas_result = st_canvas(
        stroke_width=stroke_width,
        stroke_color=stroke_color,
        background_color=bg_color,
        height=200,
        width=200,
        drawing_mode=mode,
        key="canvas",
    )
    if canvas_result.image_data is not None:
        input_img = Image.fromarray(canvas_result.image_data.astype("uint8"))

# Tab 2: Upload
with tab2:
    uploaded_file = st.file_uploader("📁 Upload a digit image...", type=["jpg", "png"])
    if uploaded_file:
        input_img = Image.open(uploaded_file).convert("RGB")

# Prediction
if input_img:
    st.image(input_img, caption="🔍 Input Image", width=150)

    if st.button("🚀 Predict Now"):
        processed = preprocess(input_img)
        prediction = model.predict(processed, verbose=0)
        digit = int(np.argmax(prediction))
        confidence = float(np.max(prediction)) * 100

        st.image(processed.reshape(28, 28), width=150, caption="🧪 Preprocessed Image")
        st.markdown(f"""
            <div class="result-box">
            🔢 Predicted Digit: <strong>{digit}</strong><br/>
            📈 Confidence: <strong>{confidence:.2f}%</strong>
            </div>
        """, unsafe_allow_html=True)