import gradio as gr
from transformers import pipeline
from PIL import Image

# Load and return a dictionary of pipelines for reuse in other applications or scripts
def load_pipeline_models():
    """
    Returns a dict of pretrained pipelines:
      - 'disease': crop disease classification pipeline
      - 'nutrient': nutrient deficiency classification pipeline
    """
    disease_pipe = pipeline(
        task="image-classification",
        model="linkanjarad/mobilenet_v2_1.0_224-plant-disease-identification",
        top_k=3
    )
    nutrient_pipe = pipeline(
        task="image-classification",
        model="nateraw/vit-base-beans",
        top_k=3
    )
    return {"disease": disease_pipe, "nutrient": nutrient_pipe}


def diagnose(image: Image.Image, models=None):
    """
    Runs inference on a PIL Image using both pipelines.

    Args:
        image (PIL.Image): Input image of a crop leaf.
        models (dict): Optional dict from load_pipeline_models(); if None, will be loaded.

    Returns:
        dict: {
            'disease': [ {'label': str, 'score': float}, ... ],
            'nutrient': [ {'label': str, 'score': float}, ... ],
            'advice': list(str)
        }
    """
    if models is None:
        models = load_pipeline_models()

    disease_results = models["disease"](image)
    nutrient_results = models["nutrient"](image)

    # Combine predictions and generate advice
    advice = []
    top_disease = disease_results[0]['label'].lower()
    top_nutrient = nutrient_results[0]['label'].lower()

    if "healthy" in top_disease:
        advice.append("No disease detected—maintain standard crop care.")
    else:
        advice.append(f"Disease detected: {disease_results[0]['label']}. Isolate and apply targeted treatment.")

    if "healthy" in top_nutrient:
        advice.append("No nutrient deficiency detected—continue regular fertilization.")
    else:
        advice.append(f"Nutrient issue: {nutrient_results[0]['label']}. Amend soil based on deficiency (e.g., add N, P, or K).")

    return {
        'disease': [ {'label': r['label'], 'score': r['score']} for r in disease_results ],
        'nutrient': [ {'label': r['label'], 'score': r['score']} for r in nutrient_results ],
        'advice': advice
    }

# Example usage
if __name__ == '__main__':
    from PIL import Image
    # Load sample image
    img = Image.open('Plants/Unhealthy_crop_1.jpg')
    models = load_pipeline_models()
    result = diagnose(img, models=models)
    print("Disease Predictions:", result['disease'])
    print("Nutrient Predictions:", result['nutrient'])
    print("Advice:", result['advice'])