from fastapi import FastAPI
from pydantic import BaseModel, field_validator
import joblib
import numpy as np

app = FastAPI()

# Load model
model = joblib.load("model.pkl")

class PredictionInput(BaseModel):
    data: list[int]

    @field_validator("data")
    def validate_length(cls, v):
        if len(v) != 17:
            raise ValueError("data must contain exactly 17 integers")
        return v

@app.post("/predict")
def predict(input_data: PredictionInput):
    X = np.array(input_data.data).reshape(1, -1)
    prediction = model.predict(X)[0]  # ambil nilai scalar
    probabilities = model.predict_proba(X)[0]  # ambil semua probabilitas untuk satu sample
    confidence = float(np.max(probabilities))  # ambil probabilitas tertinggi sebagai confidence
    
    labels = {
        0: "Normal",
        1: "Depression",
        2: "Bipolar Type-1",
        3: "Bipolar Type-2"
    }
    return {
        "prediction": int(prediction),
        "label": labels.get(int(prediction), "Unknown"),
        "confidence": round(confidence, 4)
    }


# Tambahkan baris ini jika kamu menjalankan app.py langsung!
# if __name__ == "__main__":
#     import uvicorn
#     uvicorn.run("app:app", host="0.0.0.0", port=8000)