import streamlit as st
import numpy as np
from PIL import Image
import requests
import os
import tflite_runtime.interpreter as tflite

# URL model dari Hugging Face
cnn_model_url = "https://huggingface.co/diahretnou/insectsmodel/blob/main/cnn_model.tflite"
mobilenet_model_url = "https://huggingface.co/diahretnou/insectsmodel/blob/main/convert_to_tflite.py"

@st.cache_resource
def load_model_from_url(url, filename):
    if not os.path.exists(filename):
        with st.spinner(f"Mengunduh model: {filename}..."):
            response = requests.get(url)
            with open(filename, 'wb') as f:
                f.write(response.content)
    return tf.keras.models.load_model(filename)

# Load model dari URL
cnn_model = load_model_from_url(cnn_model_url, "cnn_model.h5")
mobilenet_model = load_model_from_url(mobilenet_model_url, "mobilenet_model.h5")


# Label dan deskripsi
class_names = ['Butterfly', 'Dragonfly', 'Grasshopper', 'Ladybird', 'Mosquito']
descriptions = {
    'Grasshopper': "Grasshopper adalah serangga herbivora yang dikenal dengan kemampuan melompat jauh...",
    'Butterfly': "Butterfly adalah serangga cantik dengan sayap berwarna-warni...",
    'Dragonfly': "Dragonfly adalah serangga pemangsa yang hidup di dekat air...",
    'Ladybird': "Ladybird, atau kepik, adalah serangga kecil berwarna cerah...",
    'Mosquito': "Mosquito adalah serangga kecil yang dikenal sebagai penghisap darah..."
}

def preprocess_image(image):
    img = image.resize((150, 150))
    img = np.array(img) / 255.0
    return np.expand_dims(img, axis=0)

# UI setup
st.set_page_config(page_title="Insect Classifier", layout="wide")
st.markdown("<h1 style='text-align: center;'>🦋 Insect Classifier</h1>", unsafe_allow_html=True)
st.markdown("<h4 style='text-align: center; color: #666;'>Diah Retno Utami - 4TIB</h4>", unsafe_allow_html=True)
st.markdown("---")

# Layout 2 kolom
col1, col2 = st.columns(2)

with col1:
    st.subheader("Upload Gambar")
    uploaded_file = st.file_uploader("Pilih gambar serangga", type=["jpg", "jpeg", "png"])
    if uploaded_file is not None:
        image = Image.open(uploaded_file).convert("RGB")
        st.image(image, caption='Preview Gambar', use_column_width=True)

with col2:
    st.subheader("Hasil Prediksi")
    if uploaded_file is not None:
        img = preprocess_image(image)

        # CNN
        cnn_pred = cnn_model.predict(img)
        cnn_index = int(np.argmax(cnn_pred[0]))
        cnn_label = class_names[cnn_index]
        cnn_conf = float(np.max(cnn_pred[0]))

        # MobileNet
        mobilenet_pred = mobilenet_model.predict(img)
        mobilenet_index = int(np.argmax(mobilenet_pred[0]))
        mobilenet_label = class_names[mobilenet_index]
        mobilenet_conf = float(np.max(mobilenet_pred[0]))

        st.markdown(f"**CNN Model:** {cnn_label} ({cnn_conf*100:.2f}%)")
        st.markdown(f"**MobileNetV2 Model:** {mobilenet_label} ({mobilenet_conf*100:.2f}%)")
    else:
        st.info("Silakan upload gambar terlebih dahulu.")

# Deskripsi
if uploaded_file is not None:
    st.markdown("---")
    st.subheader("📚 Deskripsi Serangga")
    if cnn_conf >= 0.5:
        st.write(descriptions.get(cnn_label, "Deskripsi tidak tersedia."))
    else:
        st.write("Gambar tidak dapat dikenali dengan tingkat kepercayaan yang memadai.")