from fastapi import FastAPI, File, UploadFile, HTTPException
import tensorflow as tf
import numpy as np
from PIL import Image
import io
import sys
from tensorflow.keras.applications.densenet import preprocess_input

# 1. Inisialisasi Aplikasi
app = FastAPI(title="Ashoka Hipospadia Classifier API")

# 2. Load Model
print("Sedang memuat model...")
try:
    model = tf.keras.models.load_model('cnn_kfold_best_model_v2.h5')
    print("Model berhasil dimuat!")
except Exception as e:
    print(f"Error memuat model: {e}")
    sys.exit(1)  # Matikan server jika model gagal load

# Label kelas: 0 = normal, 1 = buried
class_names = ['normal', 'buried']

# 3. Fungsi Preprocessing
def prepare_image(image_bytes):
    """
    Preprocessing gambar untuk model ResNet50
    - Konversi ke RGB (3 channel)
    - Resize ke 224x224
    - Preprocessing ResNet50
    """
    try:
        img = Image.open(io.BytesIO(image_bytes))
        
        # Paksa ubah ke RGB agar PNG transparan tidak error
        img = img.convert("RGB")
        
        # Resize ke ukuran input model (224x224 untuk ResNet50)
        img = img.resize((224, 224))
        
        # Convert ke numpy array
        img_array = np.array(img)
        
        # Tambah batch dimension
        img_array = np.expand_dims(img_array, axis=0)
        
        # Preprocessing ResNet50 (HARUS sama dengan training!)
        img_array = preprocess_input(img_array)
        
        return img_array
    except Exception as e:
        print(f"Error saat memproses gambar: {e}")
        return None

# 4. Endpoint Prediksi
@app.post("/predict")
async def predict(file: UploadFile = File(...)):
    """
    Endpoint untuk prediksi gambar
    Input: File gambar (JPG, PNG, BMP)
    Output: JSON dengan class dan confidence
    """
    try:
        # Baca file gambar
        image_bytes = await file.read()
        
        # Proses gambar
        processed_image = prepare_image(image_bytes)
        
        if processed_image is None:
            raise HTTPException(status_code=400, detail="File bukan gambar yang valid")
        
        # Prediksi
        prediction = model.predict(processed_image)
        pred_value = float(prediction[0][0])
        
        # Hitung probabilitas
        # Model output: 0 = normal, 1 = buried
        prob_normal = (1 - pred_value) * 100
        prob_buried = pred_value * 100
        
        # Tentukan kelas berdasarkan threshold 0.5
        top_class_idx = 1 if pred_value > 0.5 else 0
        
        # Hasil dalam format JSON
        result = {
            "class": class_names[top_class_idx],
            "confidence": float(max(prob_normal, prob_buried)),
            "probabilities": {
                "normal": float(prob_normal),
                "buried": float(prob_buried)
            }
        }
        return result
        
    except Exception as e:
        # Cetak error ke log
        print(f"CRITICAL ERROR: {e}")
        raise HTTPException(status_code=500, detail=str(e))

# 5. Endpoint Home
@app.get("/")
def home():
    """Endpoint root untuk testing API"""
    return {
        "message": "Ashoka Hipospadia Classifier API Online! 🚀\n",
        "model": "DenseNet Binary Classification\n",
        "classes": class_names
    }

# API siap digunakan dengan uvicorn
# Jalankan dengan: uvicorn app:app --host 0.0.0.0 --port 7860