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

app = FastAPI(
    title="Iris KNN Prediction API",
    description="API for predicting Iris species using KNN model",
    version="1.0.0"
)

# Load model & class names
try:
    model, target_names = joblib.load("iris_knn.pkl")
except:
    model = None
    target_names = []

class IrisData(BaseModel):
    sepal_length: float
    sepal_width: float
    petal_length: float
    petal_width: float

@app.get("/")
def root():
    return {"message": "Iris KNN API Running! Visit /docs to test the API."}

@app.post("/predict")
def predict_iris(data: IrisData):
    if model is None:
        return {"error": "Model not found on server"}

    arr = np.array([[
        data.sepal_length,
        data.sepal_width,
        data.petal_length,
        data.petal_width
    ]])

    pred = model.predict(arr)[0]
    proba = model.predict_proba(arr)[0]

    probability_dict = {
        str(target_names[i]): float(proba[i]) for i in range(len(target_names))
    }

    return {
        "predicted_class": str(target_names[pred]),
        "input": data.dict(),
        "class_probabilities": probability_dict
    }