import gradio as gr
import torch
from transformers import ViTForImageClassification, ViTFeatureExtractor
from transformers import AutoModelForImageClassification, AutoFeatureExtractor

# Load models with error handling
try:
    model = ViTForImageClassification.from_pretrained("iamomtiwari/VITPEST")
    feature_extractor = ViTFeatureExtractor.from_pretrained("iamomtiwari/VITPEST")
    fallback_model = AutoModelForImageClassification.from_pretrained("microsoft/resnet-50")
    fallback_feature_extractor = AutoFeatureExtractor.from_pretrained("microsoft/resnet-50")
except Exception as e:
    raise gr.Error(f"Model loading failed: {str(e)}")

# Class labels and treatments (truncated for brevity)
class_labels = {
    1: {"label": "Stage Corn Common Rust", "treatment": "Apply fungicides as soon as symptoms are noticed. Practice crop rotation and remove infected plants."},
    2: {"label": "Stage Corn Gray Leaf Spot", "treatment": "Rotate crops to non-host plants, apply resistant varieties, and use fungicides as needed."},
    3: {"label": "Stage Safe Corn Healthy", "treatment": "Continue good agricultural practices: ensure proper irrigation, nutrient supply, and monitor for pests."},
    4: {"label": "Stage Corn Northern Leaf Blight", "treatment": "Remove and destroy infected plant debris, apply fungicides, and rotate crops."},
    5: {"label": "Stage Rice Brown Spot", "treatment": "Use resistant varieties, improve field drainage, and apply fungicides if necessary."},
    6: {"label": "Stage Safe Rice Healthy", "treatment": "Maintain proper irrigation, fertilization, and pest control measures."},
    7: {"label": "Stage Rice Leaf Blast", "treatment": "Use resistant varieties, apply fungicides during high-risk periods, and practice good field management."},
    8: {"label": "Stage Rice Neck Blast", "treatment": "Plant resistant varieties, improve nutrient management, and apply fungicides if symptoms appear."},
    9: {"label": "Stage Sugarcane Bacterial Blight", "treatment": "Use disease-free planting material, practice crop rotation, and destroy infected plants."},
    10: {"label": "Stage Safe Sugarcane Healthy", "treatment": "Maintain healthy soil conditions and proper irrigation."},
    11: {"label": "Stage Sugarcane Red Rot", "treatment": "Plant resistant varieties and ensure good drainage."},
    12: {"label": "Stage Wheat Brown Rust", "treatment": "Apply fungicides and practice crop rotation with non-host crops."},
    13: {"label": "Stage Safe Wheat Healthy", "treatment": "Continue with good management practices, including proper fertilization and weed control."},
    14: {"label": "Stage Wheat Yellow Rust", "treatment": "Use resistant varieties, apply fungicides, and rotate crops."}
}

labels_list = [class_labels[i]["label"] for i in range(1, 15)]
CONFIDENCE_THRESHOLD = 0.5

def predict(image):
    try:
        if not isinstance(image, torch.Tensor):
            # Convert image to tensor if needed
            inputs = feature_extractor(images=image, return_tensors="pt")
        
        with torch.no_grad():
            outputs = model(**inputs)
            logits = outputs.logits
            confidences = torch.softmax(logits, dim=-1)
            predicted_class_idx = logits.argmax(-1).item()
            confidence = confidences[0, predicted_class_idx].item()

        if confidence < CONFIDENCE_THRESHOLD:
            inputs_fallback = fallback_feature_extractor(images=image, return_tensors="pt")
            with torch.no_grad():
                outputs_fallback = fallback_model(**inputs_fallback)
                predicted_class_idx_fallback = outputs_fallback.logits.argmax(-1).item()
                fallback_label = fallback_model.config.id2label[predicted_class_idx_fallback]
            return f"Low confidence ({confidence:.2%}). Fallback prediction: {fallback_label}"

        predicted_label = labels_list[predicted_class_idx]
        treatment = class_labels[predicted_class_idx + 1]["treatment"]
        return f"Disease: {predicted_label}\n\nTreatment: {treatment}"

    except Exception as e:
        return f"Error: {str(e)}"

# Create interface with explicit types
with gr.Blocks() as demo:
    gr.Markdown("# 🌱 Crop Disease Detection")
    gr.Markdown("Upload a crop plant image to detect diseases")
    
    with gr.Row():
        image_input = gr.Image(type="pil")
        output_text = gr.Textbox()
    
    submit_btn = gr.Button("Analyze")
    submit_btn.click(
        fn=predict,
        inputs=image_input,
        outputs=output_text
    )

    gr.Examples(
        examples=[["example_corn.jpg"], ["example_wheat.jpg"]],
        inputs=image_input
    )

if __name__ == "__main__":
    demo.launch()
"""# Define class labels with treatment advice
class_labels = {
    1: {"label": "Stage Corn Common Rust", "treatment": "Apply fungicides as soon as symptoms are noticed. Practice crop rotation and remove infected plants."},
    2: {"label": "Stage Corn Gray Leaf Spot", "treatment": "Rotate crops to non-host plants, apply resistant varieties, and use fungicides as needed."},
    3: {"label": "Stage Safe Corn Healthy", "treatment": "Continue good agricultural practices: ensure proper irrigation, nutrient supply, and monitor for pests."},
    4: {"label": "Stage Corn Northern Leaf Blight", "treatment": "Remove and destroy infected plant debris, apply fungicides, and rotate crops."},
    5: {"label": "Stage Rice Brown Spot", "treatment": "Use resistant varieties, improve field drainage, and apply fungicides if necessary."},
    6: {"label": "Stage Safe Rice Healthy", "treatment": "Maintain proper irrigation, fertilization, and pest control measures."},
    7: {"label": "Stage Rice Leaf Blast", "treatment": "Use resistant varieties, apply fungicides during high-risk periods, and practice good field management."},
    8: {"label": "Stage Rice Neck Blast", "treatment": "Plant resistant varieties, improve nutrient management, and apply fungicides if symptoms appear."},
    9: {"label": "Stage Sugarcane Bacterial Blight", "treatment": "Use disease-free planting material, practice crop rotation, and destroy infected plants."},
    10: {"label": "Stage Safe Sugarcane Healthy", "treatment": "Maintain healthy soil conditions and proper irrigation."},
    11: {"label": "Stage Sugarcane Red Rot", "treatment": "Plant resistant varieties and ensure good drainage."},
    12: {"label": "Stage Wheat Brown Rust", "treatment": "Apply fungicides and practice crop rotation with non-host crops."},
    13: {"label": "Stage Safe Wheat Healthy", "treatment": "Continue with good management practices, including proper fertilization and weed control."},
    14: {"label": "Stage Wheat Yellow Rust", "treatment": "Use resistant varieties, apply fungicides, and rotate crops."}
}"""