import streamlit as st
import numpy as np
from PIL import Image
import tensorflow as tf
import tensorflow_hub as hub
from tensorflow.keras import layers
from tensorflow.keras.models import load_model

# Print versions for debugging
st.write("TensorFlow version:", tf.__version__)

# model
out_len = 10  # Replace this with the actual number of output classes

# Ensure no conflicts with 'model' or 'load_model'
model_path = 'tomato_model'

# Load your pre-trained model
try:
    VIT = load_model(model_path)
    st.write("Model loaded successfully")
except Exception as e:
    st.error(f"Error loading model: {e}")

# Define the class names
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'
]

# Define cure and prevention suggestions for each disease
disease_info = {
    'Tomato_Bacterial_spot': {
        'cure': "Remove infected plants, use copper-based fungicides.",
        'prevention': "Use disease-free seeds, practice crop rotation, avoid overhead irrigation."
    },
    'Tomato_Early_blight': {
        'cure': "Remove infected leaves, apply fungicides.",
        'prevention': "Mulch around plants, ensure good air circulation, water at the base of plants."
    },
    'Tomato_Late_blight': {
        'cure': "Remove and destroy infected plants, apply fungicides.",
        'prevention': "Plant resistant varieties, avoid overhead watering, space plants properly."
    },
    'Tomato_Leaf_Mold': {
        'cure': "Improve air circulation, apply fungicides.",
        'prevention': "Reduce humidity, avoid leaf wetness, use resistant varieties."
    },
    'Tomato_Septoria_leaf_spot': {
        'cure': "Remove infected leaves, apply fungicides.",
        'prevention': "Mulch around plants, practice crop rotation, avoid overhead watering."
    },
    'Tomato_Spider_mites_Two_spotted_spider_mite': {
        'cure': "Use insecticidal soaps or neem oil, introduce predatory mites.",
        'prevention': "Keep plants well-watered, increase humidity, use reflective mulches."
    },
    'Tomato_Target_Spot': {
        'cure': "Remove infected leaves, apply fungicides.",
        'prevention': "Improve air circulation, avoid overhead watering, practice crop rotation."
    },
    'Tomato_Tomato_Yellow_Leaf_Curl_Virus': {
        'cure': "No cure available, remove and destroy infected plants.",
        'prevention': "Use resistant varieties, control whiteflies, use reflective mulches."
    },
    'Tomato_Tomato_mosaic_virus': {
        'cure': "No cure available, remove and destroy infected plants.",
        'prevention': "Use disease-free seeds, disinfect tools, control aphids."
    },
    'Tomato_healthy': {
        'cure': "No treatment needed.",
        'prevention': "Maintain good gardening practices for overall plant health."
    }
}

# Function to load and preprocess the image
def load_and_prep_image(image):
    try:
        img = image.resize((224, 224))  # Assuming your model expects 224x224 images
        img = np.array(img) / 255.0  # Normalize the image
        img = np.expand_dims(img, axis=0)  # Add batch dimension
        return img
    except Exception as e:
        st.error(f"Error preprocessing image: {e}")
        return None

# Streamlit app
st.title("Tomato Disease Detection")
st.write("Upload an image of a tomato leaf to detect the disease.")

# File uploader
uploaded_file = st.file_uploader("Choose an image...", type="jpg")

if uploaded_file is not None:
    try:
        # Display the uploaded image
        image = Image.open(uploaded_file)
        st.image(image, caption='Uploaded Image', use_column_width=True)

        # Preprocess the image
        prepped_image = load_and_prep_image(image)

        # Ensure the model is loaded and image is preprocessed before making a prediction
        if VIT is not None and prepped_image is not None:
            # Make prediction
            prediction = VIT.predict(prepped_image)
            predicted_class = class_names[np.argmax(prediction)]

            # Display the prediction
            st.write(f"Prediction: {predicted_class}")

            # Display cure and prevention suggestions
            if predicted_class in disease_info:
                st.subheader("Cure:")
                st.write(disease_info[predicted_class]['cure'])
                st.subheader("Prevention:")
                st.write(disease_info[predicted_class]['prevention'])
            else:
                st.write("No specific cure or prevention information available for this condition.")

    except Exception as e:
        st.error(f"Error during prediction: {e}")