import torch
import torch.nn as nn
import torchvision.models as models
import torchvision.transforms as transforms
from PIL import Image
import gradio as gr

# -------------------------
# Classes (first version) - exact order from training
# -------------------------
classes = [
    'Pepper__bell___Bacterial_spot',
    'Pepper__bell___healthy',
    'Potato___Early_blight',
    'Potato___Late_blight',
    'Potato___healthy',
    '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_YellowLeaf__Curl_Virus',
    'Tomato__Tomato_mosaic_virus',
    'Tomato_healthy'
]

# Optional Hebrew translations (adjust as needed)
class_map_hebrew = {
    'Pepper__bell___Bacterial_spot': "פלפל מתוק - כתם חיידקי",
    'Pepper__bell___healthy': "פלפל מתוק - בריא",
    'Potato___Early_blight': "תפוח אדמה - ריקבון מוקדם",
    'Potato___Late_blight': "תפוח אדמה - ריקבון מאוחר",
    'Potato___healthy': "תפוח אדמה - בריא",
    'Tomato_Bacterial_spot': "עגבנייה - כתם חיידקי",
    'Tomato_Early_blight': "עגבנייה - ריקבון מוקדם",
    'Tomato_Late_blight': "עגבנייה - ריקבון מאוחר",
    'Tomato_Leaf_Mold': "עגבנייה - עובש עלה",
    'Tomato_Septoria_leaf_spot': "עגבנייה - כתם עלה Septoria",
    'Tomato_Spider_mites_Two_spotted_spider_mite': "עגבנייה - קרדית שני כתמים",
    'Tomato__Target_Spot': "עגבנייה - כתם מטרה",
    'Tomato__Tomato_YellowLeaf__Curl_Virus': "עגבנייה - צהבת עלה / וירוס התכווצות",
    'Tomato__Tomato_mosaic_virus': "עגבנייה - וירוס פסיפס",
    'Tomato_healthy': "עגבנייה - בריאה"
}

# -------------------------
# Build ResNet18 and load checkpoint
# -------------------------
model = models.resnet18(pretrained=False)
num_ftrs = model.fc.in_features
model.fc = nn.Linear(num_ftrs, len(classes))

# Load the full checkpoint (including fc)
checkpoint = torch.load("resnet18_15class.pth", map_location="cpu")
model.load_state_dict(checkpoint)
model.eval()

# -------------------------
# Image preprocessing
# -------------------------
transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406],
                         std=[0.229, 0.224, 0.225]),
])

# -------------------------
# Prediction function
# -------------------------
def predict(img):
    image = transform(img).unsqueeze(0)
    with torch.no_grad():
        outputs = model(image)
        _, predicted = torch.max(outputs, 1)
    english = classes[predicted.item()]
    hebrew = class_map_hebrew[english]
    return f"{english} ({hebrew})"

# -------------------------
# Gradio interface
# -------------------------
demo = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil"),
    outputs="text",
    title="GreenGuard - Leaf Disease Detection"
)

if __name__ == "__main__":
    demo.launch()