import gradio as gr
from tensorflow.keras.models import load_model
from tensorflow.keras.preprocessing.image import img_to_array
import numpy as np
from PIL import Image
import matplotlib.pyplot as plt
import csv
import os
from datetime import datetime

# Load model
model = load_model("waste_classification(Mobilenetv2).h5", compile=False)
class_names = ['cardboard', 'glass', 'metal', 'paper', 'plastic', 'trash']

# CSV header
csv_file = "predictions.csv"
if not os.path.exists(csv_file):
    with open(csv_file, mode='w', newline='') as file:
        writer = csv.writer(file)
        writer.writerow(["timestamp", "predicted_class", "confidence", "probabilities", "source"])

# Prediction + save to CSV
def predict_with_chart(image):
    if image is None:
        return "No image received", None

    image = image.resize((224, 224))
    img_array = img_to_array(image) / 255.0
    img_array = np.expand_dims(img_array, axis=0)

    prediction = model.predict(img_array)[0]
    pred_index = np.argmax(prediction)
    pred_label = class_names[pred_index]
    confidence = float(np.max(prediction))

    # Log prediction to CSV
    with open(csv_file, mode='a', newline='') as file:
        writer = csv.writer(file)
        writer.writerow([
            datetime.now().strftime('%Y-%m-%d %H:%M:%S'),
            pred_label,
            round(confidence, 4),
            [round(p, 4) for p in prediction.tolist()],
            "upload_or_webcam"
        ])

    # Plot class probabilities
    fig, ax = plt.subplots(figsize=(6, 4))
    ax.bar(class_names, prediction, color='skyblue')
    ax.set_ylabel('Probability')
    ax.set_ylim(0, 1)
    ax.set_title('Class Probabilities')
    plt.xticks(rotation=45)
    plt.tight_layout()

    return f"Prediction: {pred_label} ({confidence*100:.1f}%)", fig

# Gradio UI
with gr.Blocks() as demo:
    gr.Markdown("## 🗑️ Waste Classifier — Upload or Webcam\nAutomatically logs predictions to CSV.")
    with gr.Row():
        image_input = gr.Image(type="pil", label="Upload or Webcam (click camera icon)")
    with gr.Row():
        label_output = gr.Textbox(label="Predicted Class")
        plot_output = gr.Plot(label="Class Probability Chart")

    image_input.change(fn=predict_with_chart, inputs=image_input, outputs=[label_output, plot_output])

demo.launch()