import asyncio
try:
    asyncio.get_event_loop()
except RuntimeError:
    asyncio.set_event_loop(asyncio.new_event_loop())

import gradio as gr
import torch
from transformers import (
    AutoFeatureExtractor,
    AutoModelForImageClassification,
    T5Tokenizer,
    T5ForConditionalGeneration
)

# ------------------ LOAD CLASSIFIER ------------------
cls_model_name = "Aalaa/Fine_tuned_Vit_trash_classification"
feature_extractor = AutoFeatureExtractor.from_pretrained(cls_model_name)
cls_model = AutoModelForImageClassification.from_pretrained(cls_model_name)

id2label = cls_model.config.id2label


def classify_image(image):
    inputs = feature_extractor(images=image, return_tensors="pt")

    with torch.no_grad():
        outputs = cls_model(**inputs)

    logits = outputs.logits
    probs = torch.nn.functional.softmax(logits, dim=-1)[0]

    top3 = probs.topk(3).indices.tolist()
    return {id2label[i]: float(probs[i]) for i in top3}


# ------------------ LOAD FLAN-T5 ------------------
tokenizer = T5Tokenizer.from_pretrained("google/flan-t5-large")
chat_model = T5ForConditionalGeneration.from_pretrained("google/flan-t5-large")

def explain_recycling(class_label):
    prompt = f"""
    You are a waste management expert.

    The waste item is classified as: **{class_label}**
    Provide a 2-paragraph explanation including:

    1. **How to dispose of {class_label} correctly**
    2. **How {class_label} is recycled or processed**
    3. **Optional tips for reducing waste or reusing {class_label}**

    """

    inputs = tokenizer(prompt, return_tensors="pt").input_ids

    outputs = chat_model.generate(
        inputs,
        max_length=250,
        do_sample=True,
        top_p = 0.9,)

    return tokenizer.decode(outputs[0])

# ------------------ PIPELINE ------------------
def full_pipeline(image):
    predictions = classify_image(image)
    top_label = max(predictions, key=predictions.get)
    explanation = explain_recycling(top_label)
    return predictions, explanation


# ------------------ GRADIO UI ------------------
with gr.Blocks() as demo:
    gr.Markdown("<h1 style='text-align:center;'>♻️ AI Waste Classifier + Eco Advisor</h1>")

    with gr.Row():
        img_input = gr.Image(type="pil", label="Upload waste image")

    cls_output = gr.Label(num_top_classes=3, label="Classifier Prediction")

    explain_output = gr.Textbox(
        label="Detailed Recycling & Disposal Advice",
        elem_id="explainbox",
        lines=18
    )

    analyze_btn = gr.Button("Analyze", variant="primary")

    analyze_btn.click(
        full_pipeline,
        inputs=img_input,
        outputs=[cls_output, explain_output]
    )

demo.launch(
    theme=gr.themes.Soft(primary_hue="green"),
    css="#explainbox {height: 330px; font-size: 15px;}"
)