import gradio as gr
import tensorflow as tf
import numpy as np
import PIL.Image

# Class names for prediction
CLASS_NAMES = ['Alphida', 'Army worm', 'Bacterial blight', 'Healthy leaf', 'Leaf spot', 'Powdery Mildew']

# Load the pre-trained model
def load_model():
    try:
        model = tf.keras.models.load_model('cotton_plant_disease_classifier.h5')
        return model
    except Exception as e:
        print(f"Error loading model: {e}")
        return None

# Prepare the image for prediction
def prepare_image(image):
    # Resize the image
    img = image.resize((180, 180))
    # Convert to numpy array
    img_array = np.array(img)
    # Expand dimensions to create a batch
    img_array = np.expand_dims(img_array, axis=0)
    return img_array

# Prediction function
def predict_disease(model, image):
    # Prepare the image
    processed_image = prepare_image(image)
    
    # Make prediction
    predictions = model.predict(processed_image)
    score = tf.nn.softmax(predictions[0])
    
    # Get the predicted class and confidence
    predicted_class_index = np.argmax(score)
    predicted_class = CLASS_NAMES[predicted_class_index]
    confidence = 100 * np.max(score)
    
    return predicted_class, confidence

# Detailed disease information
DISEASE_INFO = {
    "Alphida": {
        "description": "Aphids are small insects that can damage cotton plants by sucking sap and spreading viruses.",
        "impact": "Low to moderate crop damage",
        "treatment": "Use insecticidal soaps, neem oil, or introduce natural predators like ladybugs"
    },
    "Army worm": {
        "description": "Army worms can cause significant damage by consuming leaf tissue, potentially reducing crop yield.",
        "impact": "High crop damage potential",
        "treatment": "Apply appropriate insecticides, practice crop rotation, maintain field hygiene"
    },
    "Bacterial blight": {
        "description": "A bacterial disease that causes lesions and wilting in cotton plants.",
        "impact": "Moderate to severe crop damage",
        "treatment": "Use copper-based bactericides, remove infected plants, practice crop rotation"
    },
    "Leaf spot": {
        "description": "A fungal disease that creates spots on leaves, potentially affecting plant health and productivity.",
        "impact": "Moderate crop damage",
        "treatment": "Apply fungicides, ensure proper plant spacing, avoid overhead irrigation"
    },
    "Powdery Mildew": {
        "description": "A fungal disease that appears as a white powdery substance on leaf surfaces.",
        "impact": "Moderate crop damage",
        "treatment": "Use sulfur-based fungicides, improve air circulation, avoid overhead watering"
    },
    "Healthy leaf": {
        "description": "The cotton plant leaf is in good health with no visible diseases or pest damage.",
        "impact": "No negative impact",
        "treatment": "Continue regular plant care and monitoring"
    }
}

# Main prediction interface
def cotton_disease_prediction(input_image):
    # Load the model (you might want to load this once globally)
    model = load_model()
    
    if model is None:
        return "Error: Model could not be loaded", None, None
    
    if input_image is None:
        return "Please upload an image", None, None
    
    try:
        # Convert to PIL Image if it's not already
        if not isinstance(input_image, PIL.Image.Image):
            input_image = PIL.Image.fromarray(input_image)
        
        # Predict disease
        predicted_class, confidence = predict_disease(model, input_image)
        
        # Get detailed information
        disease_details = DISEASE_INFO.get(predicted_class, {})
        
        # Format result
        result_text = f"Predicted Disease: {predicted_class}\n"
        result_text += f"Confidence: {confidence:.2f}%\n\n"
        result_text += f"Description: {disease_details.get('description', 'No additional information')}\n"
        result_text += f"Impact: {disease_details.get('impact', 'Not specified')}\n"
        result_text += f"Treatment: {disease_details.get('treatment', 'Consult an agricultural expert')}"
        
        return result_text, predicted_class, confidence
    
    except Exception as e:
        return f"Error in prediction: {str(e)}", None, None

# Create Gradio Interface
def create_gradio_interface():
    # Define input and output components
    image_input = gr.Image(type="pil", label="Upload Cotton Leaf Image")
    text_output = gr.Textbox(label="Prediction Results")
    disease_label = gr.Label(label="Detected Disease")
    confidence_number = gr.Number(label="Confidence Score")

    # Create the Gradio interface
    demo = gr.Interface(
        fn=cotton_disease_prediction,
        inputs=image_input,
        outputs=[text_output, disease_label, confidence_number],
        title="Cotton Plant Disease Detector",
        description="Upload a cotton plant leaf image to detect potential diseases or assess health status.",
        theme="huggingface",
        examples=[
            ["example_healthy_leaf.jpg"],
            ["example_diseased_leaf.jpg"]
        ]
    )
    
    return demo

# Launch the app
if __name__ == "__main__":
    demo = create_gradio_interface()
    demo.launch()