import mlflow 
import uvicorn
import pandas as pd 
from pydantic import BaseModel
from typing import Literal, List, Union
from fastapi import FastAPI, File, UploadFile
import joblib

# Log model from mlflow 
logged_model = 'runs:/e7b51184619c45f9b2fbb017dfe0a49f/model'

# Load model as a PyFuncModel.
loaded_model = mlflow.pyfunc.load_model(logged_model)

tags_metadata = [
    {
        "name": "Machine Learning",
        "description": "Prediction Endpoint."
    }
]

app = FastAPI(
    title="Demo Iris API",
    openapi_tags=tags_metadata
)

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

@app.get("/", tags=["Introduction Endpoints"])
async def index():
    """
    Simply returns a welcome message!
    """
    message = "Hello world! This `/` is the most simple and default endpoint. If you want to learn more, check out documentation of the api at `/docs`"
    return message


@app.post("/predict", tags=["Machine Learning"])
async def predict(predictionFeatures: PredictionFeatures):
    # Read data 
    input_data = pd.DataFrame({
            "sepal length (cm)": [predictionFeatures.sepal_length],
            "sepal width (cm)": [predictionFeatures.sepal_width],
            "petal length (cm)": [predictionFeatures.petal_length],
            "petal width (cm)": [predictionFeatures.petal_width]
        })

    prediction = loaded_model.predict(input_data)

    # Format response
    response = {"prediction": prediction.tolist()[0]}
    return response