import os
import gradio as gr
import openai
import torch
from torchvision import transforms
from PIL import Image
import requests
from io import BytesIO

# Set your OpenAI API key here or use environment variables
openai.api_key = os.getenv("OPENAI_API_KEY")

# Load a pre-trained model for leaf disease detection
# For demonstration, we'll use a generic ResNet model fine-tuned for classification
# Replace the model path with your specific model trained for leaf diseases
model = torch.hub.load('pytorch/vision:v0.10.0', 'resnet50', pretrained=True)
model.eval()

# Define image transformations
preprocess = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(
        mean=[0.485, 0.456, 0.406],  # Standard ImageNet means
        std=[0.229, 0.224, 0.225]    # Standard ImageNet stds
    )
])

# Load class labels (You should replace this with your specific disease classes)
# For demonstration, we'll use ImageNet labels
LABELS_URL = "https://s3.amazonaws.com/outcome-blog/imagenet/labels.json"
response = requests.get(LABELS_URL)
labels = {int(key): value for key, value in response.json().items()}

def detect_disease(image):
    # Preprocess the image
    img = preprocess(image).unsqueeze(0)  # Add batch dimension

    # Perform inference
    with torch.no_grad():
        outputs = model(img)
        _, predicted = torch.max(outputs, 1)
        class_id = predicted.item()
        disease = labels.get(class_id, "Unknown Disease")

    if disease == "Unknown Disease":
        return disease, "Sorry, the disease could not be identified. Please consult a local agricultural extension office."

    # Generate remedies using OpenAI's ChatGPT
    prompt = f"The following disease has been detected on a plant leaf: {disease}. Please provide detailed remedies and treatment options for this disease."

    try:
        completion = openai.ChatCompletion.create(
            model="gpt-4",
            messages=[
                {"role": "system", "content": "You are a helpful agricultural expert."},
                {"role": "user", "content": prompt}
            ],
            temperature=0.7,
            max_tokens=500
        )
        remedies = completion.choices[0].message.content.strip()
    except Exception as e:
        remedies = "An error occurred while fetching remedies. Please try again later."

    return disease, remedies

# Define the Gradio interface
iface = gr.Interface(
    fn=detect_disease,
    inputs=gr.inputs.Image(type="pil", label="Upload Leaf Image"),
    outputs=[
        gr.outputs.Textbox(label="Detected Disease"),
        gr.outputs.Textbox(label="Remedies")
    ],
    title="Leaf Disease Detector",
    description="Upload an image of a leaf, and the system will detect the disease and provide remedies."
)

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