# app.py (File ini akan diunggah ke Hugging Face Space Anda)
from flask import Flask, request, jsonify
from PIL import Image
import numpy as np
import io
import os
import tensorflow as tf # Import tensorflow di sini karena model ML akan dijalankan di sini

app = Flask(__name__)

# PATH MODEL ANDA DI HUGGING FACE SPACE
# Asumsikan 'best_model.keras' berada di root direktori Space
MODEL_PATH = "best_model.keras" 

# Inisialisasi model di luar endpoint untuk efisiensi
model = None

# Muat model saat aplikasi Flask ini dimulai
try:
    if os.path.exists(MODEL_PATH):
        model = tf.keras.models.load_model(MODEL_PATH, compile=False)
        # tf.keras.models.load_model akan otomatis mengompilasi model jika compile=True.
        # Karena kita hanya untuk inferensi, compile=False agar lebih cepat loading dan hemat memori.
        print(f"Model '{MODEL_PATH}' berhasil dimuat!")
    else:
        print(f"ERROR: File model tidak ditemukan di '{MODEL_PATH}'. API tidak akan berfungsi.")
        # Set model menjadi None agar endpoint prediksi mengembalikan error yang jelas
        model = None 
except Exception as e:
    print(f"ERROR: Gagal memuat model dari '{MODEL_PATH}': {e}. API tidak akan berfungsi.")
    model = None

# Konfigurasi Label dan Ukuran Gambar (HARUS SAMA DENGAN SAAT PELATIHAN MODEL)
CLASS_LABELS = ['battery', 'biological', 'cardboard', 'clothes', 'glass', 'metal', 'paper', 'plastic', 'shoes', 'trash']
IMG_SIZE = (150, 150) # Ukuran gambar yang diharapkan oleh model Anda

# Fungsi untuk preprocessing gambar
def preprocess_image(image_bytes):
    """
    Melakukan preprocessing gambar:
    1. Membuka gambar dari bytes.
    2. Mengonversi ke mode RGB.
    3. Mengubah ukuran (resize) gambar ke IMG_SIZE yang ditentukan.
    4. Mengonversi gambar menjadi array NumPy.
    5. Normalisasi nilai piksel (0-255 menjadi 0.0-1.0).
    6. Menambahkan dimensi batch (untuk format input model Keras/TensorFlow).
    """
    img = Image.open(io.BytesIO(image_bytes)).convert('RGB')
    img = img.resize(IMG_SIZE)
    img_array = np.array(img) / 255.0 # Normalisasi
    return np.expand_dims(img_array, axis=0) # Tambah dimensi batch

# Endpoint utama untuk mengecek API status
@app.route("/", methods=["GET"])
def home():
    """Endpoint utama untuk memverifikasi bahwa API berjalan."""
    return "TrashGu ML API is running!"

# Endpoint untuk klasifikasi gambar
@app.route("/classify", methods=["POST"])
def classify_image():
    """
    Endpoint ini menerima gambar via POST request,
    melakukan prediksi menggunakan model ML, dan mengembalikan hasilnya.
    """
    if model is None:
        # Jika model tidak berhasil dimuat saat startup
        return jsonify({"error": "Model ML belum dimuat. Mohon cek log Space Hugging Face."}), 500

    # Pastikan request memiliki file gambar
    if 'image' not in request.files:
        return jsonify({"error": "Tidak ada file gambar yang disediakan."}), 400

    image_file = request.files['image']
    image_bytes = image_file.read() # Baca file gambar sebagai bytes

    # Lakukan preprocessing gambar
    try:
        processed_image = preprocess_image(image_bytes)
    except Exception as e:
        app.logger.error(f"Gagal melakukan preprocessing gambar: {e}")
        return jsonify({"error": f"Gagal melakukan preprocessing gambar: {e}"}), 400

    # --- Bagian Deteksi Anomali (Opsional, jika ada Autoencoder) ---
    # Jika tim ML Anda memiliki model Autoencoder untuk deteksi anomali,
    # kode untuk itu akan ditambahkan di sini.
    # Contoh struktur (Anda perlu mengimplementasikan detailnya):
    is_anomaly = False
    anomaly_score = 0.0
    # if anomaly_model: # Jika anomaly_model dimuat
    #     try:
    #         reconstructed_image = anomaly_model.predict(processed_image)
    #         reconstruction_error = np.mean(np.square(processed_image - reconstructed_image))
    #         anomaly_score = float(reconstruction_error)
    #         # Tentukan threshold anomali Anda
    #         if reconstruction_error > 0.05: # Contoh threshold
    #             is_anomaly = True
    #     except Exception as e:
    #         print(f"Warning: Deteksi anomali gagal: {e}")
    #         # Lanjutkan tanpa deteksi anomali jika terjadi error
    # Jika terdeteksi anomali, berikan respons anomali
    if is_anomaly:
        return jsonify({
            "prediction": "TIDAK_DIKETAHUI", # Label khusus jika anomali
            "confidence": 0.0,
            "is_anomaly": True,
            "anomaly_score": anomaly_score
        }), 200 # Status 200 OK karena deteksi anomali berhasil

    # --- Klasifikasi Sampah Menggunakan Model ---
    try:
        predictions = model.predict(processed_image)
        # Ambil array prediksi pertama jika outputnya nested (misal: [[0.1, 0.9]])
        pred_array = predictions[0] if isinstance(predictions, np.ndarray) and predictions.ndim > 1 else predictions 
        
        # Dapatkan indeks kelas dengan probabilitas tertinggi
        predicted_class_index = np.argmax(pred_array)
        # Dapatkan label kelas dari CLASS_LABELS
        predicted_label = CLASS_LABELS[predicted_class_index]
        # Dapatkan nilai kepercayaan (confidence)
        confidence = float(pred_array[predicted_class_index])

        # Mengembalikan hasil dalam format JSON
        return jsonify({
            "prediction": predicted_label, # Contoh: 'plastic', 'biological'
            "confidence": confidence,     # Contoh: 0.98
            "is_anomaly": False,          # Default ke False jika tidak ada deteksi anomali atau tidak terdeteksi
            "anomaly_score": 0.0          # Default ke 0.0
        }), 200

    except Exception as e:
        app.logger.error(f"Gagal melakukan inferensi model: {e}")
        return jsonify({"error": f"Gagal melakukan inferensi model: {e}"}), 500

# Bagian ini hanya berjalan saat file dijalankan langsung (bukan saat di-deploy oleh WSGI)
if __name__ == "__main__":
    # Hugging Face Spaces secara otomatis menjalankan aplikasi di port yang sesuai.
    # Anda tidak perlu menentukan host/port secara eksplisit jika menggunakan SDK seperti 'Flask'.
    # Baris di bawah ini lebih untuk pengujian lokal.
    app.run(debug=True, host='0.0.0.0', port=int(os.environ.get("PORT", 5000)))