import os
import numpy as np
from flask import Flask, request, jsonify
from tensorflow.keras.models import load_model
from tensorflow.keras.preprocessing.image import img_to_array
from PIL import Image
from rembg import remove
from flask_cors import CORS

app = Flask(__name__)
CORS(app)

# 1. Load Model
# Pastikan file model terbaru sudah kamu download dan ganti namanya menjadi ini
MODEL_PATH = 'train_part4.keras' 
model = load_model(MODEL_PATH)

# Pastikan urutan kelas sesuai dengan test_generator.class_indices
# Tadi di Colab urutannya: honey, natural, wash (bukan washed)
classes = ['honey', 'natural', 'wash']

# 2. Fungsi Preprocessing Robust (Sesuai eksperimen terakhir di Colab)
def preprocess_image(input_img):
    # A. Hapus Background
    output_rgba = remove(input_img)
    
    # B. Auto-Crop ke Bounding Box (Fokus ke biji kopi saja)
    bbox = output_rgba.getbbox()
    if bbox:
        output_rgba = output_rgba.crop(bbox)
        
    # C. Center Padding (Membuat kanvas hitam persegi)
    max_dim = max(output_rgba.size)
    black_bg = Image.new("RGB", (max_dim, max_dim), (0, 0, 0))
    
    # Hitung posisi agar biji kopi di tengah
    paste_x = (max_dim - output_rgba.size[0]) // 2
    paste_y = (max_dim - output_rgba.size[1]) // 2
    
    # Tempelkan gambar transparan ke latar hitam
    black_bg.paste(output_rgba, (paste_x, paste_y), mask=output_rgba.split()[3])
    
    # D. Resize ke 224x224 (Input EfficientNetB0)
    final_img = black_bg.resize((224, 224))
    
    # E. Konversi ke Array
    img_array = img_to_array(final_img)
    img_array = np.expand_dims(img_array, axis=0)
    return img_array

@app.route('/predict', methods=['POST'])
def predict():
    if 'image' not in request.files:
        return jsonify({"status": "ERROR", "message": "File gambar tidak ditemukan"}), 400
        
    image_file = request.files['image']
    try:
        # Load gambar asli sebagai RGBA agar rembg bekerja maksimal
        img = Image.open(image_file.stream).convert("RGBA")
        
        # Jalankan Preprocessing Robust
        processed_img = preprocess_image(img)
        
        # Prediksi
        preds = model.predict(processed_img)[0]
        confidence = float(np.max(preds))
        predicted_class = classes[np.argmax(preds)]
        
        # Hitung Entropy (Mengukur tingkat kebingungan model)
        entropy = -np.sum(preds * np.log(preds + 1e-9))
        
        prob_details = "Analisis Probabilitas Model:\n"
        for i, cls_name in enumerate(classes):
            prob_details += f"- {cls_name.capitalize()}: {round(float(preds[i]) * 100, 2)}%\n"
        prob_details = prob_details.strip()
        
        # --- LOGIKA PENYARINGAN (Threshold & Entropy) ---
        
        # 1. Jika entropy tinggi (Model bingung parah)
        if entropy > 0.85:
            return jsonify({
                "status": "DITOLAK",
                "pesan": f"Sistem mendeteksi ketidakjelasan. Pastikan objek adalah biji kopi tunggal dengan pencahayaan cukup.\n\n{prob_details}",
                "entropy_score": round(float(entropy), 4)
            }), 200

        # 2. Jika Keyakinan Rendah (Di bawah 70%)
        if confidence < 0.70:
            return jsonify({
                "status": "TIDAK YAKIN",
                "label": predicted_class,
                "confidence": str(round(confidence * 100, 2)),
                "pesan": f"Model menduga ini proses {predicted_class}, namun tingkat keyakinan rendah.\n\n{prob_details}"
            }), 200

        # 3. Lolos Verifikasi (Status SUKSES)
        return jsonify({
            "status": "SUKSES",
            "label": predicted_class,
            "confidence": str(round(confidence * 100, 2)),
            "pesan": f"Biji kopi teridentifikasi sebagai proses {predicted_class}.\n\n{prob_details}"
        }), 200

    except Exception as e:
        return jsonify({"status": "ERROR", "message": str(e)}), 500

if __name__ == '__main__':
    # Pastikan port sesuai dengan yang dibuka di firewall server/local kamu
    app.run(host='0.0.0.0', port=5001, debug=False)