import streamlit as st
import numpy as np
from PIL import Image
from tensorflow.keras.models import load_model
from streamlit_drawable_canvas import st_canvas


ms = st.session_state
if "themes" not in ms: 
  ms.themes = {"current_theme": "light",
                    "refreshed": True,
                    
                    "light": {"theme.base": "dark",
                              "theme.backgroundColor": "black",
                              "theme.primaryColor": "#c98bdb",
                              "theme.secondaryBackgroundColor": "#5591f5",
                              "theme.textColor": "white",
                              "theme.textColor": "white",
                              "button_face": "🌜"},

                    "dark":  {"theme.base": "light",
                              "theme.backgroundColor": "white",
                              "theme.primaryColor": "#5591f5",
                              "theme.secondaryBackgroundColor": "#82E1D7",
                              "theme.textColor": "#0a1464",
                              "button_face": "🌞"},
                    }
  

def ChangeTheme():
  previous_theme = ms.themes["current_theme"]
  tdict = ms.themes["light"] if ms.themes["current_theme"] == "light" else ms.themes["dark"]
  for vkey, vval in tdict.items(): 
    if vkey.startswith("theme"): st._config.set_option(vkey, vval)

  ms.themes["refreshed"] = False
  if previous_theme == "dark": ms.themes["current_theme"] = "light"
  elif previous_theme == "light": ms.themes["current_theme"] = "dark"


btn_face = ms.themes["light"]["button_face"] if ms.themes["current_theme"] == "light" else ms.themes["dark"]["button_face"]
st.button(btn_face, on_click=ChangeTheme)

if ms.themes["refreshed"] == False:
  ms.themes["refreshed"] = True
  st.rerun()

# Function to preprocess the image
def preprocess_image(image, target_size):
    if image.mode != "RGB":
        image = image.convert("RGB")
    image = image.resize(target_size)
    image = np.expand_dims(image, axis=0)
    return image

# Function to predict the digit
def predict_digit(model, image):
    processed_image = preprocess_image(image, (200, 200))  # Match your model's input size
    prediction = model.predict(processed_image)
    return np.argmax(prediction), np.max(prediction)

# Load your trained model
model = load_model("last_burmese_Digit_recognizer_model.h5")

# Streamlit app
st.title("Burmese Digit Recognizer")

# Upload image file or draw
st.markdown("## Upload an Image or Draw")
col1, col2 = st.columns(2)

with col1:
    file = st.file_uploader("Upload Here", type=['png', 'jpg', 'jpeg'])

with col2:
    # Drawable canvas
    canvas_result = st_canvas(
        fill_color="rgba(255, 165, 0, 0.3)",  # Drawing parameters
        stroke_width=3,
        stroke_color="#ffffff",
        background_color="#000000",
        background_image=None if file else st.session_state.get("background", None),
        update_streamlit=True,
        width=400,
        height=400,
        drawing_mode="freedraw",
        key="canvas",
    )

image = None  # Initialize image variable

# Process uploaded image or drawing
if file is not None:
    image = Image.open(file)  # Read image with PIL
elif canvas_result.image_data is not None:
    image = Image.fromarray(np.array(canvas_result.image_data, dtype=np.uint8)).convert('RGB')

if image is not None:
    st.image(image, caption='Uploaded Image')  # Display the uploaded/drawn image

    # Predict the digit
    digit, confidence = predict_digit(model, image)
    st.write(f"Predicted Digit: {digit} with confidence {confidence:.2f}")
else:
    st.write("Please upload an image or use the canvas to draw.")