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 import MobileNetV2
from tensorflow.keras.applications.mobilenet_v2 import preprocess_input
from tensorflow.keras.layers import GlobalAveragePooling2D, Dense, Dropout
from tensorflow.keras.models import Model

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

# ===============================
# 2. Build Model (SAMA DENGAN TRAINING)
# ===============================
def build_model():
    base_model = MobileNetV2(
        input_shape=(224, 224, 3),
        include_top=False,
        weights="imagenet"
    )
    base_model.trainable = False

    x = base_model.output
    x = GlobalAveragePooling2D()(x)
    x = Dropout(0.2)(x)
    output = Dense(1, activation="sigmoid")(x)

    model = Model(inputs=base_model.input, outputs=output)
    return model


# ===============================
# 3. Load Weights
# ===============================
print("Sedang memuat model...")
try:
    model = build_model()
    model.load_weights("sultan_model.weights.h5")
    print("Model berhasil dimuat!")
except Exception as e:
    print(f"Error memuat model: {e}")
    sys.exit(1)

# Label kelas
class_names = ["normal", "buried"]

# ===============================
# 4. Preprocessing Image
# ===============================
def prepare_image(image_bytes):
    try:
        img = Image.open(io.BytesIO(image_bytes))
        img = img.convert("RGB")
        img = img.resize((224, 224))

        img_array = np.array(img)
        img_array = np.expand_dims(img_array, axis=0)
        img_array = preprocess_input(img_array)

        return img_array
    except Exception as e:
        print(f"Error preprocessing image: {e}")
        return None


# ===============================
# 5. Endpoint Prediksi
# ===============================
@app.post("/predict")
async def predict(file: UploadFile = File(...)):
    try:
        image_bytes = await file.read()
        processed_image = prepare_image(image_bytes)

        if processed_image is None:
            raise HTTPException(status_code=400, detail="Invalid image file")

        prediction = model.predict(processed_image)
        pred_value = float(prediction[0][0])

        prob_normal = (1 - pred_value) * 100
        prob_buried = pred_value * 100

        predicted_class = "buried" if pred_value > 0.5 else "normal"

        return {
            "class": predicted_class,
            "confidence": max(prob_normal, prob_buried),
            "probabilities": {
                "normal": prob_normal,
                "buried": prob_buried
            }
        }

    except Exception as e:
        print(f"CRITICAL ERROR: {e}")
        raise HTTPException(status_code=500, detail=str(e))


# ===============================
# 6. Endpoint Root
# ===============================
@app.get("/")
def home():
    return {
        "message": "Ashoka Hipospadia Classifier API Online 🚀",
        "model": "MobileNetV2 Binary Classification",
        "classes": class_names
    }