# streamlit_app.py
import streamlit as st
from fastai.vision.all import *
import matplotlib.pyplot as plt
import matplotlib.image as mpimg

# Function to get the label from the file name
def get_label(file_name):
    return file_name.split('-')[0]

# Function to prepare data (similar to your code)
def prepare_data(food_path, label_a, label_b):
    for img in get_image_files(food_path):
        if label_a in str(img):
            img.rename(f"{img.parent}/{label_a}-{img.name}")
        elif label_b in str(img):
            img.rename(f"{img.parent}/{label_b}-{img.name}")
        else:
            os.remove(img)

# Function to train the model
def train_model(food_path, label_func):
    dls = ImageDataLoaders.from_name_func(
        food_path, get_image_files(food_path), valid_pct=0.2, seed=420,
        label_func=label_func, item_tfms=Resize(230)
    )

    learn = cnn_learner(dls, resnet34, metrics=error_rate, pretrained=True)
    learn.fine_tune(epochs=1)

    return learn

# Streamlit app
def main():
    st.title("Food Classifier Streamlit App")

    # Sidebar options
    options = ["Train Model", "Upload Image", "Test Random Images", "Confusion Matrix"]
    choice = st.sidebar.selectbox("Choose an option", options)

    if choice == "Train Model":
        st.subheader("Training the Model")
        food_path = Path("~/.fastai/data/food-101/food-101").expanduser()
        if not food_path.exists():
            try:
                food_path = untar_data(URLs.FOOD)
            except FileExistsError:
                st.warning("Data directory already exists. Skipping download.")
            label_a = st.text_input("Enter label A:", "samosa")
            label_b = st.text_input("Enter label B:", "hot_and_sour_soup")

            if st.button("Train Model"):
                prepare_data(food_path, label_a, label_b)
                learn = train_model(food_path, get_label)
                st.session_state.model = learn  # Save the model to session state
                st.success("Model trained successfully!")

    elif choice == "Upload Image":
        st.subheader("Upload Your Own Images")
        if "model" not in st.session_state:
            st.warning("Please train the model first.")
        else:
            uploaded_files = st.file_uploader("Choose images", type=["jpg", "jpeg", "png"], accept_multiple_files=True)

            if uploaded_files:
                for img in uploaded_files:
                    img = PILImage.create(img)
                    label, _, probs = st.session_state.model.predict(img)

                    st.image(img, caption=f"This is a {label}.")
                    st.write(f"{label}: {probs[1].item():.6f}")
                    st.write(f"{label}: {probs[0].item():.6f}")

    elif choice == "Test Random Images":
        st.subheader("Test Using Images in Dataset")
        if "model" not in st.session_state:
            st.warning("Please train the model first.")
        else:
            for i in range(0, 5):  # Change 5 to the number of images you want to display
                random_index = random.randint(0, len(get_image_files(food_path)) - 1)
                img_path = get_image_files(food_path)[random_index]
                img = mpimg.imread(img_path)
                label, _, probs = st.session_state.model.predict(img)

                st.image(img, caption=f"Predicted label: {label}")

    elif choice == "Confusion Matrix":
        st.subheader("Confusion Matrix")
        if "model" not in st.session_state:
            st.warning("Please train the model first.")
        else:
            interp = ClassificationInterpretation.from_learner(st.session_state.model)
            st.pyplot(interp.plot_confusion_matrix())

# Run the Streamlit app
if __name__ == "__main__":
    main()