import streamlit as st
import numpy as np
import cv2
import pickle
from pathlib import Path

# ------------------------------
# Paths
# ------------------------------
SRC_PATH = Path(__file__).parent  # points to src/

# ------------------------------
# Load models
# ------------------------------
with open(SRC_PATH / "pca_model.pkl", "rb") as f:
    pca = pickle.load(f)

with open(SRC_PATH / "svc_model.pkl", "rb") as f:
    svc = pickle.load(f)

with open(SRC_PATH / "scaler.pkl", "rb") as f:
    scaler = pickle.load(f)

# ------------------------------
# Label mapping
# ------------------------------
label_map = {
    0: "fresh apple",
    1: "fresh banana",
    2: "fresh orange",
    3: "rotten apple",
    4: "rotten banana",
    5: "rotten orange"
}

# ------------------------------
# Streamlit UI
# ------------------------------
st.title("🍎 Fruit Image Classification")

# ---- Example images ----
st.subheader("Example Images")
example_images = [
    "Fresh Apple.jpg",
    "Fresh Banana.jpg",
    "Fresh Orange.jpg",
    "Rotten Apple.jpg",
    "Rotten Banana.jpg",
    "Rotten Orange.jpg"
]

cols = st.columns(3)
for i, img_file in enumerate(example_images):
    img_path = SRC_PATH / img_file
    if img_path.exists():
        img = cv2.imread(str(img_path))
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        cols[i % 3].image(img, caption=img_file.split(".")[0], width=150)
    else:
        st.warning(f"Image '{img_file}' not found.")

st.markdown("**Classes:** fresh apple, fresh banana, fresh orange, rotten apple, rotten banana, rotten orange")

# ------------------------------
# File uploader
# ------------------------------
uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])

if uploaded_file is not None:
    # Read uploaded file as OpenCV image
    file_bytes = np.asarray(bytearray(uploaded_file.read()), dtype=np.uint8)
    img = cv2.imdecode(file_bytes, cv2.IMREAD_COLOR)

    # Convert BGR → RGB for display
    img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    st.image(img_rgb, caption="Uploaded Image", use_container_width=True)

    # ------------------------------
    # Preprocess & predict
    # ------------------------------
    img_resized = cv2.resize(img, (64, 64))
    img_flat = img_resized.flatten().reshape(1, -1)
    img_scaled = scaler.transform(img_flat)
    img_pca = pca.transform(img_scaled)
    pred = svc.predict(img_pca)[0]
    label = label_map.get(pred, "Unknown")

    st.success(f"Predicted Class: {label}")