import streamlit as st
import tensorflow as tf
from PIL import Image
import numpy as np
import json
import os
import pandas as pd

st.set_page_config(page_title="Food Image Classification", layout="centered")

@st.cache_resource
def load_model_and_metadata():
    """
    Load the saved Keras model (best_model.keras) and metadata.json
    from the current directory.
    """
    cwd = os.getcwd()
    # 1. Load model file
    model_file = os.path.join(cwd, "best_model.keras")
    if not os.path.exists(model_file):
        st.error(f"Model file not found at '{model_file}'. Ensure best_model.keras is present.")
        return None, None, None, None
    try:
        model = tf.keras.models.load_model(model_file)
    except Exception as e:
        st.error(f"Failed to load model from {model_file}: {e}")
        return None, None, None, None

    # 2. Load metadata.json at root
    meta_path = os.path.join(cwd, "metadata.json")
    if not os.path.exists(meta_path):
        st.error(f"metadata.json not found at '{meta_path}'.")
        return model, None, None, None
    try:
        with open(meta_path, "r") as f:
            metadata = json.load(f)
    except Exception as e:
        st.error(f"Failed to read metadata.json: {e}")
        return model, None, None, None

    class_names = metadata.get("class_names")
    IMG_HEIGHT = metadata.get("IMG_HEIGHT")
    IMG_WIDTH = metadata.get("IMG_WIDTH")
    if class_names is None or IMG_HEIGHT is None or IMG_WIDTH is None:
        st.error("metadata.json must contain 'class_names', 'IMG_HEIGHT', and 'IMG_WIDTH'")
        return model, None, None, None

    return model, class_names, IMG_HEIGHT, IMG_WIDTH


model, class_names, IMG_HEIGHT, IMG_WIDTH = load_model_and_metadata()
if model is None or class_names is None:
    st.stop()



st.title("Food Image Classification")
st.write("Upload a food image (Bread / Soup / Vegetable-Fruit) to classify.")

with st.sidebar:
    st.header("Instructions")
    st.write(
        f"""
        - Upload a JPG/PNG image of food.
        - The model expects images resized to {IMG_HEIGHT}×{IMG_WIDTH}.
        - The model classes: {', '.join(class_names)}.
        - For best results, upload a clear image of a single food item.
        """
    )


st.subheader("Image Prediction")
uploaded_file = st.file_uploader("Choose an image...", type=["jpg","jpeg","png"], key="single")

if uploaded_file is not None:
    try:
        image = Image.open(uploaded_file).convert("RGB")
    except Exception as e:
        st.error(f"Cannot open image: {e}")
        image = None
    if image:
        st.image(image, caption="Uploaded Image", use_column_width=True)

        # Preprocess: resize and scale to [0,1]
        img_resized = image.resize((IMG_WIDTH, IMG_HEIGHT))
        img_array = np.array(img_resized).astype("float32") / 255.0
        input_tensor = np.expand_dims(img_array, axis=0)  # shape (1, H, W, 3)

        # Predict
        preds = model.predict(input_tensor, verbose=0)
        pred_idx = int(np.argmax(preds[0]))
        pred_class = class_names[pred_idx]
        confidence = float(preds[0][pred_idx])
        st.write(f"**Prediction:** {pred_class}  \n**Confidence:** {confidence:.3f}")

        # Show full probability distribution
        prob_df = pd.DataFrame({
            "class": class_names,
            "probability": preds[0]
        })
        st.bar_chart(data=prob_df.set_index("class"))

        # Download result as CSV
        result_df = pd.DataFrame([{
            "filename": uploaded_file.name,
            "predicted_class": pred_class,
            "confidence": confidence
        }])

        csv = result_df.to_csv(index=False).encode('utf-8')

        st.download_button(
            label="Download prediction as CSV",
            data=csv,
            file_name="prediction.csv",
            mime="text/csv"
        )