from fastapi import FastAPI, HTTPException, Header
from pydantic import BaseModel
from typing import List, Dict, Any, Optional
import base64

app = FastAPI()


from typing import List
import numpy as np
import joblib
import pandas as pd
import os
import traceback

# Load the real model
MODEL_PATH = os.path.join(os.path.dirname(__file__), "iris_knn_pipeline.pkl")
model = None
load_error = None

try:
    if os.path.exists(MODEL_PATH):
        model = joblib.load(MODEL_PATH)
    else:
        load_error = f"Model file not found at {MODEL_PATH}"
except Exception as e:
    load_error = f"Error loading model: {str(e)}\n{traceback.format_exc()}"

def predict_iris(features: List[float]) -> tuple:
    if load_error:
        raise Exception(f"Model not loaded: {load_error}")
    
    # Feature names expected by the scikit-learn pipeline
    FEATURE_NAMES = [
        "sepal length (cm)",
        "sepal width (cm)",
        "petal length (cm)",
        "petal width (cm)"
    ]
    
    # Convert to DataFrame with correct column names
    df = pd.DataFrame([features], columns=FEATURE_NAMES)
    
    pred = int(model.predict(df)[0])
    probs = model.predict_proba(df)[0].tolist()
    
    return pred, probs

CLASS_NAMES = ["setosa", "versicolor", "virginica"]


class ArrayRequest(BaseModel):
    features: List[float]

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

@app.get("/")
def root():
    return {"message": "Iris API - Format 4 - Array Wrapper"}


@app.post("/predict")
async def predict(request: ArrayRequest):
    pred, probs = predict_iris(request.features)
    return [
        {
            "prediction": pred,
            "probabilities": probs
        }
    ]


if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)