import streamlit as st
import tensorflow as tf
import numpy as np
from PIL import Image
import io
import logging

# Set up logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Load the pre-trained MNIST model
@st.cache_resource
def load_model():
    try:
        model = tf.keras.models.load_model('mnist_cnn.h5')
        logger.info("MNIST model loaded successfully")
        return model
    except Exception as e:
        logger.error(f"Error loading model: {e}")
        st.error("Failed to load the model. Please check the model file.")
        return None

# Preprocess the uploaded image
def preprocess_image(image):
    try:
        # Convert to grayscale
        img = image.convert('L')
        # Resize to 28x28 (MNIST model input size)
        img = img.resize((28, 28), Image.Resampling.LANCZOS)
        # Convert to numpy array and normalize
        img_array = np.array(img)
        # Ensure the image is inverted if necessary (MNIST expects white digits on black background)
        img_array = 255 - img_array  # Invert colors
        img_array = img_array / 255.0  # Normalize to [0, 1]
        # Reshape for model input (1, 28, 28, 1)
        img_array = img_array.reshape(1, 28, 28, 1)
        logger.info("Image preprocessed successfully")
        return img_array
    except Exception as e:
        logger.error(f"Error preprocessing image: {e}")
        st.error("Failed to preprocess the image. Please ensure it's a valid image.")
        return None

# Streamlit app
st.title("AutoWeightLogger - Number Detection")
st.write("Upload an image containing a single handwritten digit to detect the number.")

# File uploader
uploaded_file = st.file_uploader("Choose an image...", type=["png", "jpg", "jpeg"])

if uploaded_file is not None:
    try:
        # Display the uploaded image
        image = Image.open(uploaded_file)
        st.image(image, caption="Uploaded Image", use_column_width=True)

        # Preprocess the image
        processed_image = preprocess_image(image)
        if processed_image is None:
            st.stop()

        # Load the model
        model = load_model()
        if model is None:
            st.stop()

        # Make prediction
        with st.spinner("Detecting number..."):
            prediction = model.predict(processed_image)
            predicted_digit = np.argmax(prediction, axis=1)[0]
            confidence = np.max(prediction) * 100

        # Display result
        st.success(f"Detected Number: {predicted_digit}")
        st.write(f"Confidence: {confidence:.2f}%")

        # Provide feedback if confidence is low
        if confidence < 70:
            st.warning("Low confidence in prediction. Please ensure the image contains a clear, single handwritten digit.")

    except Exception as e:
        logger.error(f"Error processing image: {e}")
        st.error("An error occurred while processing the image. Please try again with a different image.")
else:
    st.info("Please upload an image to proceed.")

# Instructions for users
st.markdown("""
### Instructions
1. Upload an image containing a single handwritten digit (0-9).
2. Ensure the digit is clear, centered, and on a plain background for best results.
3. The model expects white digits on a black background, similar to MNIST dataset images.
""")