import gradio as gr
import tensorflow as tf
import numpy as np
import json
from PIL import Image

# Load model and class mapping
model = tf.keras.models.load_model("cropguard_model.h5")

with open("class_indices.json", "r") as f:
    class_indices = json.load(f)

idx_to_class = {v: k for k, v in class_indices.items()}

# Human-readable labels and treatment tips
disease_info = {
    "Pepper__bell___Bacterial_spot": {
        "name": "Bell Pepper — Bacterial Spot",
        "tip": "Remove infected leaves, avoid overhead watering, apply copper-based bactericide."
    },
    "Pepper__bell___healthy": {
        "name": "Bell Pepper — Healthy",
        "tip": "No action needed. Maintain good watering and spacing practices."
    },
    "Potato___Early_blight": {
        "name": "Potato — Early Blight",
        "tip": "Remove affected foliage, rotate crops yearly, apply fungicide if severe."
    },
    "Potato___Late_blight": {
        "name": "Potato — Late Blight",
        "tip": "Highly destructive — remove infected plants immediately, apply fungicide, avoid wet foliage."
    },
    "Potato___healthy": {
        "name": "Potato — Healthy",
        "tip": "No action needed. Continue regular monitoring."
    },
    "Tomato_Bacterial_spot": {
        "name": "Tomato — Bacterial Spot",
        "tip": "Avoid overhead watering, remove infected leaves, apply copper-based spray."
    },
    "Tomato_Early_blight": {
        "name": "Tomato — Early Blight",
        "tip": "Remove lower infected leaves, mulch soil, apply fungicide preventatively."
    },
    "Tomato_Late_blight": {
        "name": "Tomato — Late Blight",
        "tip": "Act fast — highly contagious. Remove and destroy infected plants, apply fungicide."
    },
    "Tomato_Leaf_Mold": {
        "name": "Tomato — Leaf Mold",
        "tip": "Improve air circulation, reduce humidity, apply fungicide if persistent."
    },
    "Tomato_Septoria_leaf_spot": {
        "name": "Tomato — Septoria Leaf Spot",
        "tip": "Remove infected lower leaves, avoid wetting foliage, rotate crops."
    },
    "Tomato_Spider_mites_Two_spotted_spider_mite": {
        "name": "Tomato — Spider Mites",
        "tip": "Spray with insecticidal soap or neem oil, increase humidity around plants."
    },
    "Tomato__Target_Spot": {
        "name": "Tomato — Target Spot",
        "tip": "Remove infected debris, apply fungicide, improve air circulation."
    },
    "Tomato__Tomato_YellowLeaf__Curl_Virus": {
        "name": "Tomato — Yellow Leaf Curl Virus",
        "tip": "No cure — remove infected plants to prevent spread, control whiteflies (the vector)."
    },
    "Tomato__Tomato_mosaic_virus": {
        "name": "Tomato — Mosaic Virus",
        "tip": "No cure — remove and destroy infected plants, disinfect tools between use."
    },
    "Tomato_healthy": {
        "name": "Tomato — Healthy",
        "tip": "No action needed. Continue regular care and monitoring."
    },
}

def predict_disease(img):
    if img is None:
        return "Please upload a leaf image.", ""

    img = img.resize((224, 224))
    img_array = np.array(img) / 255.0
    img_array = np.expand_dims(img_array, axis=0)

    preds = model.predict(img_array, verbose=0)[0]
    pred_idx = np.argmax(preds)
    confidence = round(float(preds[pred_idx]) * 100, 1)

    raw_label = idx_to_class[pred_idx]
    info = disease_info.get(raw_label, {"name": raw_label, "tip": "No info available."})

    result = f"{info['name']}  ({confidence}% confidence)"
    tip = f"💡 Recommendation: {info['tip']}"

    # Top 3 predictions
    top3_idx = np.argsort(preds)[::-1][:3]
    breakdown = "Top 3 predictions:\n"
    for idx in top3_idx:
        lbl = disease_info.get(idx_to_class[idx], {"name": idx_to_class[idx]})["name"]
        pct = round(float(preds[idx]) * 100, 1)
        breakdown += f"  • {lbl}: {pct}%\n"

    return result, tip, breakdown

with gr.Blocks(title="CropGuard") as demo:
    gr.Markdown("""
    # 🌱 CropGuard — Crop Disease Detector
    ### Upload a photo of a Tomato, Potato, or Bell Pepper leaf to detect disease
    *Built by Samuel Yaula Dutse | MobileNetV2 Transfer Learning | 93% Accuracy*
    """)

    with gr.Row():
        with gr.Column():
            image_input = gr.Image(type="pil", label="Upload Leaf Image")
            submit_btn = gr.Button("Diagnose", variant="primary")

        with gr.Column():
            result_output = gr.Textbox(label="Diagnosis", interactive=False)
            tip_output = gr.Textbox(label="Recommendation", interactive=False)
            breakdown_output = gr.Textbox(label="Confidence Breakdown", lines=5, interactive=False)

    submit_btn.click(
        fn=predict_disease,
        inputs=[image_input],
        outputs=[result_output, tip_output, breakdown_output]
    )

demo.launch()