from pathlib import Path
import platform
import pathlib

import gradio as gr
from fastai.vision.all import PILImage, load_learner

if platform.system() == "Windows":
    pathlib.PosixPath = pathlib.WindowsPath


# Load model once when app starts
MODEL_PATH = Path(__file__).parent / "models" / "waste_classifier_v2.pkl"
learn = load_learner(MODEL_PATH)

# Waste bin mapping
BIN_INFO = {
    "plastic": "Plastic Recycling",
    "glass": "Glass Recycling",
    "metal": "Metal Recycling",
    "paper": "Paper Recycling",
    "cardboard": "Paper/Cardboard Recycling",
    "organic": "Organic/Compost",
    "textile": "Textile Collection",
    "e-waste": "Electronic Waste Collection",
}


def predict(image):
    img = PILImage.create(image)

    pred, pred_idx, probs = learn.predict(img)
    confidence = float(probs[pred_idx])

    bin_name = BIN_INFO[str(pred)]

    return (
        f"## {pred.title()} ({confidence:.1%})",
        f"**Recommended disposal:** {bin_name}",
    )


demo = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil", label="Upload a waste image"),
    outputs=[
        gr.Markdown(label="Prediction"),
        gr.Markdown(label="Disposal Recommendation"),
    ],
    title="♻️ AI Trash Sorter",
    description="Upload an image of waste and the AI will identify it and recommend the correct trash bin.",
)

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