import gradio as gr
import numpy as np
from PIL import Image
from tensorflow.keras.models import load_model
from tensorflow.keras.applications.efficientnet import preprocess_input as efficientnet_preprocess
import pandas as pd
import os
from datetime import datetime

# Waste categories and disposal advice
waste_advice = {
    "PAPER": {
        "category": "Paper",
        "advice": "Recyclable: Clean paper, newspapers, magazines, cardboard. Remove staples and plastic windows. Avoid: Wax-coated paper, tissues, heavily soiled paper.",
        "bin": "Blue Recycling Bin"
    },
    "PLASTIC": {
        "category": "Plastic",
        "advice": "Check recycling numbers (1-7). Commonly accepted: bottles, containers. Clean containers before recycling. Take plastic bags to store drop-off points.",
        "bin": "Yellow Recycling Bin"
    },
    "GLASS": {
        "category": "Glass",
        "advice": "Rinse bottles and jars. Remove metal lids. Sort by colour if required.",
        "bin": "Green Recycling Bin"
    },
    "METAL": {
        "category": "Metal",
        "advice": "Aluminium cans are highly recyclable. Remove labels and rinse steel cans. Take larger items to scrapyards.",
        "bin": "Red Recycling Bin"
    },
    "ORGANIC": {
        "category": "Organic",
        "advice": "Compost fruit/vegetable scraps, coffee grounds, eggshells. Avoid meat and dairy in home compost.",
        "bin": "Brown Compost Bin"
    },
    "BATTERIES": {
        "category": "Battery",
        "advice": "Take to designated collection points at supermarkets or electronics stores. Never put in regular rubbish due to toxic materials.",
        "bin": "Special Collection Point"
    },
    "E-WASTE": {
        "category": "E-Waste",
        "advice": "Take to manufacturer take-back programmes or electronics retailers. Wipe personal data before disposal. Contains valuable materials.",
        "bin": "E-Waste Collection Centre"
    },
    "CLOTHES": {
        "category": "Clothes",
        "advice": "Donate clean, wearable items. Use textile recycling bins for worn-out clothes. Many retailers offer take-back programmes.",
        "bin": "Donation Centre/Textile Bin"
    },
    "LIGHT-BULBS": {
        "category": "Light Bulbs",
        "advice": "Do not place in general recycling. Take CFL and fluorescent bulbs to special collection sites.",
        "bin": "Special Disposal Site"
    }
}


class_names = ['BATTERIES', 'CLOTHES', 'E-WASTE', 'GLASS', 'LIGHT-BULBS', 'METAL', 'ORGANIC', 'PAPER', 'PLASTIC']

# Load model
try:
    model = load_model("model_efnet.keras")
except OSError:
    model = None
    print("⚠️ Model file not found. Please ensure 'model_efnet.keras' exists in the directory.")

    

img_size = (224, 224)
log_file = "prediction_log.csv"
if not os.path.exists(log_file):
    pd.DataFrame(columns=["Timestamp", "Prediction", "Confidence", "Feedback"]).to_csv(log_file, index=False)

def predict(image):
    if model is None:
        return "Model not loaded", "", None, None, {}

    img = image.convert("RGB").resize(img_size)
    if not isinstance(image, Image.Image):
        return "Invalid image", "", None, None, {}

    img_array = efficientnet_preprocess(np.array(img))
    img_array = np.expand_dims(img_array, axis=0)
    prediction = model.predict(img_array, verbose=0)[0]
    predicted_index = np.argmax(prediction)
    predicted_class = class_names[predicted_index]
    confidence = prediction[predicted_index]

    advice_info = waste_advice.get(predicted_class, {
        "category": predicted_class,
        "advice": "No disposal advice available.",
        "bin": "Not specified"
    })

    formatted_advice = f"""
    <h3><b>Category:</b> {advice_info['category']}</h3>
    <p><b>Advice:</b> {advice_info['advice']}</p>
    <p><b>Bin:</b> {advice_info['bin']}</p>
    """

    result = f"{predicted_class} ({confidence:.2%})"
    probabilities = {class_names[i]: float(prediction[i]) for i in range(len(class_names))}
    return result, formatted_advice, predicted_class, confidence, probabilities

def save_feedback(predicted_class, confidence, feedback):
    timestamp = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
    if confidence is None:
        confidence_value = ""
    else:
        confidence_value = round(confidence, 2)

    df = pd.read_csv(log_file)
    new_row = pd.DataFrame({
        "Timestamp": [timestamp],
        "Prediction": [predicted_class],
        "Confidence": [round(confidence, 2)],
        "Feedback": [feedback]
    })
    df = pd.concat([df, new_row], ignore_index=True)
    df.to_csv(log_file, index=False)
    
    return(
        gr.update(visible=True),               # feedback_ack
        gr.update(value="", visible=False),   # feedback_box
        gr.update(value=None)                 # thumbs
    )

def handle_feedback(feedback):
    if feedback == "👎 Wrong":
        return (
            gr.update(visible=True),              # Show textbox
            gr.update(visible=True),              # Show submit button
            gr.update(visible=False)              # Hide feedback_ack
        )
    return (
        gr.update(visible=False),                 # Hide textbox
        gr.update(visible=False),                 # Hide submit button
        gr.update(value="Thank you", visible=True)  # Show feedback_ack with message
    )

def clear_all():
    return None, "", "", None, None, {}, None, gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)

def load_images_from_folder():
    folder_path = "sample images"  # Your folder name
    image_files = []
    
    # Get all image files from the folder
    for filename in os.listdir(folder_path):
        if filename.lower().endswith(('.png', '.jpg', '.jpeg', '.gif', '.bmp')):
            image_path = os.path.join(folder_path, filename)
            image_files.append(image_path)
    
    return image_files

# Gradio UI
with gr.Blocks() as demo:
    gr.Markdown("""
    # ♻️ Smart Waste  Classifier
    Upload a waste image or Use webcam to get a predicted category and disposal advice.
    """)
    with gr.Row():
        image_input = gr.Image(label="Upload or Use Webcam", sources=["upload", "webcam"], type="pil")
        with gr.Column():
            result = gr.Label(label="Prediction")
            advice_output = gr.HTML(label="Waste Disposal Advice")

    pred_class = gr.State()
    pred_conf = gr.State()
    probabilities = gr.State()

    with gr.Row():
        predict_button = gr.Button("Detect waste")
        clear_btn = gr.Button("Clear")

    predict_button.click(
        fn=predict,
        inputs=image_input,
        outputs=[result, advice_output, pred_class, pred_conf, probabilities]
    )

    gr.Markdown("### How accurate was the prediction?")
    with gr.Row():
        thumbs = gr.Radio(["👍 Correct", "👎 Wrong"], label="Prediction Status")
        feedback_box = gr.Textbox(lines=2, placeholder="What should the model have predicted?", visible=False, label="Your Feedback")
        submit_feedback = gr.Button("Submit Feedback", visible=False)
        feedback_ack = gr.Textbox(visible=False, label="✅ Feedback saved. Thank you!", interactive=False)

    clear_btn.click(
        fn=clear_all,
        outputs=[image_input, result, advice_output, pred_class, pred_conf, probabilities, thumbs, feedback_box, submit_feedback, feedback_ack]
    )
    thumbs.change(
        fn=handle_feedback,
        inputs=[thumbs],
        outputs=[feedback_box, submit_feedback, feedback_ack]
    )
    submit_feedback.click(
        fn=save_feedback,
        inputs=[pred_class, pred_conf, feedback_box],
        outputs=[feedback_ack, feedback_box, thumbs]
    )
   # Inside Gradio UI
    gr.Markdown("#### ✨Image Gallery: Test with these sample images")
    
    # Display images in a gallery
    gallery = gr.Gallery(
        value=load_images_from_folder(),
        label="Sample Images",
        show_label=True,
        elem_id="gallery",
        columns=6,
        rows=1,
        height="auto"
    )

    
demo.launch()