# -*- coding: utf-8 -*-

import pandas as pd
from pycaret.regression import load_model, predict_model
from fastapi import FastAPI
from pydantic import BaseModel
import uvicorn

# Create the app
app = FastAPI()

# Load trained Pipeline
model = load_model("lr_api")

# Create input/output pydantic models
class InputModel(BaseModel):
    rownames: int
    year: int
    violent: float
    murder: float
    prisoners: int
    afam: float
    cauc: float
    male: float
    population: float
    income: float
    density: float
    state: str
    law: str

class OutputModel(BaseModel):
    prediction: float

# Define predict function
@app.post("/predict", response_model=OutputModel)
def predict(data: InputModel):
    data = pd.DataFrame([data.dict()])
    predictions = predict_model(model, data=data)
    return {"prediction": predictions["prediction_label"].iloc[0]}


#if __name__ == "__main__":
#    uvicorn.run(app, host="127.0.0.1", port=8000)