import streamlit as st
import os
import requests
import tensorflow as tf
from tensorflow.keras.models import load_model
import numpy as np
from PIL import Image

# 🔧 Configure the Streamlit page
st.set_page_config(page_title="Leukemia Subtype Detection", layout="centered")
st.markdown("<h1 style='text-align: center; color: #d6336c;'>🎗️ Leukemia Subtype Detection</h1>", unsafe_allow_html=True)
st.markdown("<p style='text-align: center; font-size:18px;'>Upload a blood smear image and detect the leukemia subtype using deep learning models.</p>", unsafe_allow_html=True)

# 🔗 Hugging Face Base URL (confirmed valid)
HF_BASE = "https://huggingface.co/GarimaSharma75/Leukemia_Subtype/resolve/main/"
SAVE_DIR = "models"

# 📌 Class labels
CLASS_NAMES = ["Early Pre-B ALL", "Pre-B ALL", "Pro-B ALL", "Healthy"]

# 🔢 Model files hosted on Hugging Face
model_files = {
    "DenseNet121": "DenseNet121_model.keras",
    "MobileNetV2": "MobileNetV2_model.keras",
    "VGG16": "VGG16_model.keras",
    "CustomCNN": "CustomCNN_model.keras"
}

# 📥 Download the model file from Hugging Face if it's not saved locally
def download_if_not_exists(filename):
    os.makedirs(SAVE_DIR, exist_ok=True)
    filepath = os.path.join(SAVE_DIR, filename)
    if not os.path.exists(filepath):
        url = HF_BASE + filename
        with st.spinner(f"📥 Downloading `{filename}`..."):
            response = requests.get(url)
            if response.status_code == 200:
                with open(filepath, "wb") as f:
                    f.write(response.content)
            else:
                st.error(f"❌ Failed to download {filename} from Hugging Face.")
                st.stop()
    return filepath

# 🧼 Image preprocessing
def preprocess(img):
    img = img.resize((224, 224))
    img_array = np.array(img) / 255.0
    return np.expand_dims(img_array, axis=0)

# 📍 Sidebar: Model selection and info
with st.sidebar:
    st.markdown("## 🧠 Select Model")
    selected_model = st.selectbox("Choose one model to run", list(model_files.keys()))
    st.markdown("### ℹ️ About the Models")
    st.info("""
    • **DenseNet121** – Deep CNN with dense connections  
    • **MobileNetV2** – Lightweight CNN  
    • **VGG16** – Classic 16-layer CNN  
    • **CustomCNN** – Custom-built architecture
    """)
    st.markdown("---")
    st.markdown("Made by **Garima Sharma** 💖")

# 📤 Upload image
uploaded_file = st.file_uploader("📤 Upload a blood smear image (JPG/PNG)", type=["jpg", "jpeg", "png"])

if uploaded_file:
    img = Image.open(uploaded_file).convert("RGB")
    st.image(img, caption="Uploaded Image", use_container_width=True)

    if st.button("🔍 Run Detection"):
        with st.spinner("⏳ Please wait while the model is downloading and predicting..."):
            input_data = preprocess(img)
            model_path = download_if_not_exists(model_files[selected_model])
            model = load_model(model_path)
            preds = model.predict(input_data)
            pred_class = CLASS_NAMES[np.argmax(preds)]
            prob = np.max(preds) * 100

        st.success(f"✅ **{selected_model}** predicts: **{pred_class}** with `{prob:.2f}%` confidence")
else:
    st.warning("📎 Please upload an image to get started.")