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

# === CONFIG ===
MODEL_PATH = 'trained_model/asl_model.h5'
IMG_SIZE = 64
CLASS_NAMES = [chr(i) for i in range(65, 91)]  # A-Z

# Load model once
@st.cache_resource(show_spinner=False)
def load_model():
    return tf.keras.models.load_model(MODEL_PATH)

model = load_model()

# === UI Header ===
st.set_page_config(page_title="ASL Recognition", page_icon="🧠", layout="centered")
st.markdown("<h1 style='text-align: center;'>🧠 ASL Alphabet Recognition</h1>", unsafe_allow_html=True)
st.markdown("<p style='text-align: center;'>Upload a hand gesture image and get instant letter prediction.</p>", unsafe_allow_html=True)
st.divider()

# === Helper Functions ===
def preprocess_image(image: Image.Image):
    img = image.convert("RGB")
    img = img.resize((IMG_SIZE, IMG_SIZE))
    img = np.array(img) / 255.0
    img = np.expand_dims(img, axis=0)
    return img

def predict(img: Image.Image):
    processed = preprocess_image(img)
    preds = model.predict(processed)
    class_idx = np.argmax(preds)
    confidence = preds[0][class_idx]
    return CLASS_NAMES[class_idx], confidence

# === Upload UI ===
uploaded_file = st.file_uploader("📁 Upload a hand gesture image", type=['png', 'jpg', 'jpeg'])

if uploaded_file:
    col1, col2 = st.columns([1, 2])
    with col1:
        img = Image.open(uploaded_file)
        st.image(img, caption="📷 Uploaded Image", use_column_width=True)
    with col2:
        st.write("### 🔍 Prediction")
        label, confidence = predict(img)
        st.success(f"Predicted Letter: **:blue[{label}]**")
        st.metric(label="Confidence Score", value=f"{confidence * 100:.2f}%", delta=None)

        # Optional: show full probabilities as a horizontal bar chart
        preds = model.predict(preprocess_image(img))[0]
        top_indices = np.argsort(preds)[::-1][:5]
        st.write("#### 🔢 Top 5 Predictions")
        for i in top_indices:
            st.progress(float(preds[i]), text=f"{CLASS_NAMES[i]}: {preds[i]*100:.2f}%")
else:
    st.info("📸 Upload a clear image showing a single hand gesture on a plain background.")

# === Footer ===
st.divider()
st.markdown("<small style='text-align:center; display:block;'>Developed with ❤️ using TensorFlow & Streamlit</small>", unsafe_allow_html=True)