import gradio as gr
import tensorflow as tf
import json
import tensorflow_hub as hub
from PIL import Image

# Ensure KerasLayer is recognized when loading the model
tf.keras.utils.get_custom_objects().update({'KerasLayer': hub.KerasLayer})

# Paths to your model and label files
class_file_path = './labels.json'
model_file_path = './model.h5'

# Load the model
model = tf.keras.models.load_model(model_file_path)

def load_breeds(file_path=class_file_path):
    with open(file_path, 'r') as file:
        return json.load(file)

labels = load_breeds()

# Format disease list as comma-separated string, nicely readable for UI
def format_label(label):
    label = label.replace('__', ' ')
    label = label.replace('_', ' ')
    return label.title()

disease_list_str = ", ".join([format_label(label) for label in labels])

# Organic treatments dictionary with keys matching the raw label strings exactly
organic_treatments = {
    "Maize Rust": "Spray neem oil or copper fungicide. Use resistant maize varieties and remove infected plant debris.",
    "Maize fall armyworm": "Handpick larvae, use Bacillus thuringiensis (Bt) sprays, and encourage natural predators like birds and parasitic wasps.",
    "Maize grasshoper": "Introduce natural predators such as birds, use neem-based insecticides, and practice crop rotation.",
    "Maize healthy": "No treatment needed. Maintain good agricultural practices and crop hygiene.",
    "Maize leaf beetle": "Use neem oil sprays, release beneficial insects like ladybugs, and remove affected leaves.",
    "Maize leaf blight": "Apply copper-based fungicides, use resistant varieties, and avoid overhead irrigation to reduce leaf wetness.",
    "Maize leaf spot": "Remove and destroy infected leaves, use copper fungicides, and ensure proper spacing for air circulation.",
    "Maize streak virus": "Control insect vectors like leafhoppers with neem insecticide, plant resistant varieties, and remove infected plants.",
    "Tomato_Bacterial_spot": "Spray copper-based bactericides, remove infected plant material, and avoid overhead watering.",
    "Tomato_Early_blight": "Use copper fungicides, remove and destroy infected leaves, and rotate crops.",
    "Tomato_Late_blight": "Apply organic fungicides like copper or bicarbonate sprays, remove infected plants, and avoid wetting foliage.",
    "Tomato_Leaf_Mold": "Ensure good air circulation, avoid overhead watering, and apply neem or copper fungicides.",
    "Tomato_Septoria_leaf_spot": "Remove infected leaves, use copper fungicides, and maintain proper plant spacing.",
    "Tomato_Spider_mites_Two_spotted_spider_mite": "Spray insecticidal soap or neem oil, introduce predatory mites, and regularly hose plants to remove mites.",
    "Tomato__Target_Spot": "Remove affected leaves, use copper fungicides, and practice crop rotation.",
    "Tomato__Tomato_YellowLeaf__Curl_Virus": "Control whitefly vectors with neem insecticide, remove infected plants, and use resistant varieties.",
    "Tomato__Tomato_mosaic_virus": "Use virus-free seeds, disinfect tools, remove infected plants, and practice crop rotation.",
    "Tomato_healthy": "No treatment needed. Maintain proper watering, good air circulation, and balanced fertilization."
}

def process_image(image, img_size=224):
    img_array = tf.keras.preprocessing.image.img_to_array(image)
    img_array = tf.image.resize(img_array, [img_size, img_size]) / 255.0
    return img_array

def predict_breed(image):
    try:
        if image is None:
            return "❌ No image uploaded. Please upload a Maize or Tomato leaf image.", {}, ""

        img_array = process_image(image)
        img_array = tf.expand_dims(img_array, axis=0)

        predictions = model.predict(img_array)[0]
        top3_indices = predictions.argsort()[-3:][::-1]

        top_pred_class_raw = labels[top3_indices[0]]  # raw key
        top_pred_class = format_label(top_pred_class_raw)

        # Check for valid crops in the display label
        if "Maize" not in top_pred_class and "Tomato" not in top_pred_class:
            return "❌ This model only supports Maize and Tomato leaf images.", {}, ""

        output_lines = ["🌿 **Top 3 Predictions:**"]
        confidence_scores = {}

        for i in top3_indices:
            class_name = format_label(labels[i])
            confidence = predictions[i] * 100
            output_lines.append(f"- **{class_name}**: {confidence:.2f}%")
            confidence_scores[class_name] = float(f"{confidence:.2f}")

        treatment = organic_treatments.get(top_pred_class_raw, "No organic treatment available.")

        treatment_text = f"🌱 **Organic Treatment for {top_pred_class}:**\n\n{treatment}"

        return "\n".join(output_lines), confidence_scores, treatment_text

    except Exception as e:
        print("Prediction Error:", e)
        return "❌ File not supported. Please upload a valid image file (JPEG/PNG).", {}, ""

with gr.Blocks() as demo:
    gr.Markdown("# 🌿 Hares: Maize & Tomato Disease Classifier")
    gr.Markdown(f"### Supported Diseases:\n\n{disease_list_str}")

    image_input = gr.Image(type="pil", label="Upload Maize or Tomato Leaf Image")
    prediction_text = gr.Markdown(label="Prediction")
    confidence_bar = gr.JSON(label="Confidence Scores")
    treatment_text = gr.Markdown(label="Organic Treatment")

    image_input.change(fn=predict_breed, inputs=image_input, outputs=[prediction_text, confidence_bar, treatment_text])

demo.launch(share=True)