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

# -------------------------------
# Load pre-trained image classifier
# -------------------------------
# Small and fast model for demo
try:
    image_classifier = pipeline("image-classification", model="google/vit-base-patch16-224")
except Exception as e:
    print("⚠️ Could not load image classifier:", e)
    image_classifier = None

# -------------------------------
# Waste categories
# -------------------------------
compostable = [
    "vegetable", "vegetables", "tomato", "onion", "potato", "carrot",
    "fruit", "apple", "banana", "orange", "mango", "food", "leaves", "cardboard", "paper"
]

recyclable = ["plastic", "bottle", "can", "glass", "metal", "aluminum", "tin", "carton"]

harmful = ["syringe", "battery", "medical", "medicine", "chemical", "paint", "electronics", "toxic"]

# -------------------------------
# Text classification
# -------------------------------
def classify_text(item):
    item = item.lower()
    if any(word in item for word in compostable):
        return "✅ Compostable Waste (Green Bin)"
    elif any(word in item for word in recyclable):
        return "♻️ Recyclable Waste (Blue Bin)"
    elif any(word in item for word in harmful):
        return "⚠️ Harmful/Non-Decomposable Waste (Red Bin)"
    else:
        return "🚮 Unknown Waste (Grey Bin)"

# -------------------------------
# Image classification
# -------------------------------
def classify_image(image_path):
    if image_classifier is None:
        return "⚠️ Image classifier not available."
    try:
        img = Image.open(image_path)
        preds = image_classifier(img, top_k=5)
        labels = [p["label"].lower() for p in preds]

        if any(word in labels for word in compostable):
            return "✅ Compostable Waste (Green Bin)"
        elif any(word in labels for word in recyclable):
            return "♻️ Recyclable Waste (Blue Bin)"
        elif any(word in labels for word in harmful):
            return "⚠️ Harmful/Non-Decomposable Waste (Red Bin)"
        else:
            return "🚮 Unknown Waste (Grey Bin)"
    except Exception as e:
        return f"⚠️ Error classifying image: {e}"

# -------------------------------
# Main classifier
# -------------------------------
def classify_waste(input_type, image_input, text_input):
    if input_type == "Image Upload / Webcam" and image_input is not None:
        return classify_image(image_input)
    elif input_type == "Text Description" and text_input:
        return classify_text(text_input)
    else:
        return "❌ Please provide input for the selected option."

# -------------------------------
# Gradio UI
# -------------------------------
with gr.Blocks() as demo:
    gr.Markdown("# ♻️ EcoSort: Smart Waste Classifier")
    gr.Markdown("Upload a photo (or use webcam) **OR** type a description to see which bin the item belongs to!")

    input_type = gr.Radio(
        ["Image Upload / Webcam", "Text Description"],
        label="Select input type",
        value="Text Description"
    )

    with gr.Row():
        image_input = gr.Image(type="filepath", label="Upload or capture photo", sources=["upload", "webcam"])
        text_input = gr.Textbox(label="Type the waste item description")

    output = gr.Textbox(label="Classification Result")

    classify_btn = gr.Button("Classify")
    classify_btn.click(
        fn=classify_waste,
        inputs=[input_type, image_input, text_input],
        outputs=output
    )

# -------------------------------
# Launch
# -------------------------------
if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)