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

app = FastAPI()

# -------- Load models once at startup --------
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
MODEL_DIR = os.path.join(BASE_DIR, "models")

stockout_model = joblib.load(
    os.path.join(MODEL_DIR, "restaurant_stockout_classifier.joblib")
)

wastage_model = joblib.load(
    os.path.join(MODEL_DIR, "restaurant_wastage_regressor.joblib")
)

# -------- Request schema --------
class PredictRequest(BaseModel):
    features: list[float]

# -------- Health check --------
@app.get("/")
def root():
    return {
        "status": "ok",
        "message": "ProjectY Classifier + Regressor API is running"
    }

# -------- Stockout classifier --------
@app.post("/predict/stockout")
def predict_stockout(req: PredictRequest):
    X = np.array([req.features])  # shape: (1, n_features)
    prediction = stockout_model.predict(X)[0]

    response = {
        "prediction": int(prediction) if isinstance(prediction, (int, np.integer)) else float(prediction)
    }

    # Optional probabilities (if supported)
    if hasattr(stockout_model, "predict_proba"):
        response["probabilities"] = stockout_model.predict_proba(X)[0].tolist()

    return response

# -------- Wastage regressor --------
@app.post("/predict/wastage")
def predict_wastage(req: PredictRequest):
    X = np.array([req.features])
    prediction = wastage_model.predict(X)[0]

    return {
        "prediction": float(prediction)
    }