import os
import cv2
import numpy as np
import joblib
import skimage.feature as ft
from flask import Flask, request, jsonify
from flask_cors import CORS
from rembg import remove
import matplotlib.pyplot as plt

# Inisialisasi Flask
app = Flask(__name__)
CORS(app)  # Mengizinkan akses dari Flutter

# Path model dan scaler
model_path = r'C:\Strawberry\classifragise\flask.backend\strawberry_freshness_knn.pkl'
scaler_path = r'C:\Strawberry\classifragise\flask.backend\scaler_knn.pkl'

# Load model & scaler
knn_model = joblib.load(model_path)
scaler = joblib.load(scaler_path)

# Folder untuk menyimpan gambar sementara
UPLOAD_FOLDER = "uploads"
os.makedirs(UPLOAD_FOLDER, exist_ok=True)

def remove_background(input_path, output_path):
    """Menghapus latar belakang gambar"""
    try:
        with open(input_path, "rb") as input_file:
            img_data = input_file.read()
            output_data = remove(img_data)

        with open(output_path, "wb") as output_file:
            output_file.write(output_data)
        
        print(f"✅ Background dihapus untuk {input_path}")
    except Exception as e:
        print(f"❌ Gagal menghapus background: {e}")

def refine_segmentation(image_path):
    """Refine hasil segmentasi setelah penghapusan background"""
    img = cv2.imread(image_path, cv2.IMREAD_UNCHANGED)
    
    if img is None:
        print("❌ Gagal membaca gambar untuk refining segmentation")
        return

    # Konversi ke grayscale
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

    # Coba threshold lebih lembut
    _, binary = cv2.threshold(gray, 30, 255, cv2.THRESH_BINARY)

    # Operasi morfologi untuk menghilangkan noise
    kernel = np.ones((2,2), np.uint8)
    binary = cv2.morphologyEx(binary, cv2.MORPH_CLOSE, kernel)

    # Simpan hasil refinemen
    cv2.imwrite(image_path, binary)
    print("✅ Segmentasi background diperbaiki!")

def extract_features(image_path):
    """Ekstraksi fitur RGB, GLCM, dan Laplacian dari gambar"""
    img = cv2.imread(image_path)

    if img is None:
        raise ValueError(f"Gambar tidak dapat dibaca: {image_path}")

    # Konversi ke grayscale
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

    # Rata-rata nilai RGB
    mean_r = np.mean(img[:, :, 2])
    mean_g = np.mean(img[:, :, 1])
    mean_b = np.mean(img[:, :, 0])

    # GLCM features
    glcm = ft.graycomatrix(gray, distances=[1], angles=[0], levels=256, symmetric=True, normed=True)
    contrast = ft.graycoprops(glcm, 'contrast')[0, 0]
    homogeneity = ft.graycoprops(glcm, 'homogeneity')[0, 0]
    energy = ft.graycoprops(glcm, 'energy')[0, 0]
    correlation = ft.graycoprops(glcm, 'correlation')[0, 0]

    # Laplacian features
    laplacian = cv2.Laplacian(gray, cv2.CV_64F)
    mean_laplacian = np.mean(laplacian)
    var_laplacian = np.var(laplacian)

    return np.array([mean_r, mean_g, mean_b, contrast, homogeneity, energy, correlation, mean_laplacian, var_laplacian]).reshape(1, -1)

@app.route("/")
def home():
    return jsonify({"message": "API Flask untuk Klasifikasi Stroberi Aktif!"})

@app.route("/predict", methods=["POST"])
def predict():
    if knn_model is None or scaler is None:
        return jsonify({"error": "Model atau Scaler belum dimuat"}), 500

    if "file" not in request.files:
        return jsonify({"error": "Tidak ada file gambar yang dikirim"}), 400

    file = request.files["file"]

    if file.filename == "":
        return jsonify({"error": "Nama file kosong"}), 400

    # Simpan gambar sementara
    file_path = os.path.join(UPLOAD_FOLDER, file.filename)
    file.save(file_path)

    # Path hasil setelah background dihapus
    file_nobg = file_path.replace(".jpg", "_nobg.png")

    # Hapus background & refining
    remove_background(file_path, file_nobg)
    refine_segmentation(file_nobg)

    # **Tambahkan visualisasi hasil penghapusan background**
    try:
        img = cv2.imread(file_nobg)
        if img is not None:
            plt.imshow(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
            plt.axis("off")
            plt.title("Cek Hasil Segmen Sebelum Ekstraksi Fitur")
            plt.show()
        else:
            print("❌ Gambar hasil penghapusan background tidak ditemukan!")
    except Exception as e:
        print(f"❌ Gagal menampilkan gambar: {e}")

    # Ekstraksi fitur
    try:
        features = extract_features(file_nobg)
        print(f"✅ Fitur berhasil diekstraksi: {features}")

        # Normalisasi fitur
        features_scaled = scaler.transform(features)
        print(f"✅ Fitur setelah normalisasi: {features_scaled}")

        # Prediksi
        prediction = knn_model.predict(features_scaled)
        result = "Segar ✅" if prediction[0] == 1 else "Tidak Segar ❌"
        return jsonify({"prediction": result})

    except Exception as e:
        return jsonify({"error": f"Gagal memproses gambar: {str(e)}"}), 500

if __name__ == "__main__":
    app.run(host="0.0.0.0", port=5000, debug=True)
