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

MODEL_PATH = Path(__file__).resolve().parent / "EfficientNetV2B0_best.h5"

# Jika Anda menggunakan label lain, sesuaikan list ini dengan urutan output model Anda.
CLASS_NAMES = [
    "Tomato_Bacterial_spot",
    "Tomato_Early_blight",
    "Tomato_Late_blight",
    "Tomato_Leaf_Mold",
    "Tomato_Septoria_leaf_spot",
    "Tomato_Spider_mites_Two_spotted_spider_mite",
    "Tomato_Target_Spot",
    "Tomato_Tomato_Yellow_Leaf_Curl_Virus",
    "Tomato_Tomato_mosaic_virus",
    "Tomato_Healthy",
]

@st.cache_resource
def load_model(path: Path):
    if not path.exists():
        raise FileNotFoundError(f"Model tidak ditemukan: {path}")
    return tf.keras.models.load_model(str(path))

def preprocess_image(_image: Image.Image, target_size=(224, 224)) -> np.ndarray:
    image = _image.convert("RGB")
    image = image.resize(target_size)
    image_array = np.array(image).astype(np.float32)
    image_array = tf.keras.applications.efficientnet_v2.preprocess_input(image_array)
    image_array = np.expand_dims(image_array, axis=0)
    return image_array

def predict_image(model, _image: Image.Image):
    img_tensor = preprocess_image(_image)
    preds = model.predict(img_tensor)
    return preds[0]


def main():
    st.set_page_config(
        page_title="Tomato Leaf Disease Detector",
        page_icon="🍅",
        layout="centered",
    )

    st.title("Deteksi Penyakit Daun Tomat")
    st.write(
        "Unggah gambar daun tomat dan model akan memprediksi jenis penyakit (atau sehat) pada tanaman tersebut."
    )

    uploaded_file = st.file_uploader(
        "Pilih gambar daun tomat", type=["png", "jpg", "jpeg"]
    )

    if uploaded_file is None:
        st.info("Silakan unggah gambar untuk melakukan prediksi.")
        return

    try:
        image = Image.open(uploaded_file)
    except Exception as e:
        st.error(f"Gagal membaca gambar: {e}")
        return

    st.image(image, caption="Gambar input", use_column_width=True)

    try:
        model = load_model(MODEL_PATH)
    except Exception as e:
        st.error(f"Gagal memuat model: {e}")
        return

    with st.spinner("Menjalankan prediksi..."):
        preds = predict_image(model, image)

    num_classes = preds.shape[-1] if preds.ndim > 0 else 1
    if num_classes != len(CLASS_NAMES):
        class_names = [f"Class {i}" for i in range(num_classes)]
    else:
        class_names = CLASS_NAMES

    top_idx = int(np.argmax(preds))
    confidence = float(np.max(preds))
    predicted_label = class_names[top_idx]

    st.markdown("---")
    st.subheader("Hasil Prediksi")
    st.write(f"**Prediksi:** {predicted_label}")
    st.write(f"**Kepercayaan:** {confidence * 100:.2f}%")

    if len(preds) == len(class_names):
        st.subheader("Distribusi Skor Kelas")
        result = {class_names[i]: float(preds[i]) for i in range(len(preds))}
        st.bar_chart(result)

    st.write(
        "\n*Catatan: Sesuaikan `CLASS_NAMES` di `app.py` jika urutan kelas dalam model Anda berbeda.*"
    )


if __name__ == "__main__":
    main()