# Source: content/notes/libraries/mlops-and-serving.md
# Independent CPU example; use the curriculum environment.
# See /notes/ml/#example-environment or /notes/deep-learning/#example-environment.

from contextlib import asynccontextmanager
import numpy as np
from fastapi import FastAPI, HTTPException
from fastapi.testclient import TestClient
from pydantic import BaseModel, Field
from sklearn.linear_model import LogisticRegression

class Request(BaseModel):
    f0: float = Field(..., allow_inf_nan=False)
    f1: float = Field(..., allow_inf_nan=False)
    class Config:
        extra = "forbid"

def create_app(loader):
    @asynccontextmanager
    async def lifespan(app):
        model, version = loader()
        warm = model.predict_proba(np.zeros((1, 2)))
        if warm.shape != (1, 2) or not np.isfinite(warm).all():
            raise RuntimeError("warmup failed")
        app.state.model, app.state.version = model, version
        app.state.ready = True
        try:
            yield
        finally:
            app.state.ready = False

    app = FastAPI(lifespan=lifespan)
    app.state.ready = False

    @app.get("/health")
    def health():
        return {"ok": True}

    @app.get("/ready")
    def ready():
        if not app.state.ready:
            raise HTTPException(503, "not ready")
        return {"ok": True}

    @app.post("/predict")
    def predict(request: Request):
        if not app.state.ready:
            raise HTTPException(503, "not ready")
        positive = np.flatnonzero(app.state.model.classes_ == 1).item()
        score = app.state.model.predict_proba([[request.f0, request.f1]])[0, positive]
        return {"score": float(score), "model_version": app.state.version}
    return app

X = np.array([[-2., 0], [-1., 1], [1., -1], [2., 0]])
model = LogisticRegression(random_state=0).fit(X, [0, 0, 1, 1])
app = create_app(lambda: (model, "fixture-version-7"))
assert TestClient(app).get("/ready").status_code == 503
with TestClient(app) as client:
    assert client.get("/ready").status_code == 200
    result = client.post("/predict", json={"f1": 0., "f0": 2.}).json()
    assert result["model_version"] == "fixture-version-7"
    assert np.isclose(result["score"], model.predict_proba([[2., 0.]])[0, 1])
    assert client.post("/predict", json={"f0": 1.}).status_code == 422
    assert client.post("/predict", json={"f0": 1., "f1": 0., "secret": 3}).status_code == 422
assert not app.state.ready
print("readiness status, fixed schema, class probability and immutable version passed")