from utils import Preprocess, MissingValue
import pickle
from fastapi import FastAPI, Request
import pandas as pd
import numpy as np
import onnxruntime as ort


# Load pre-trained model, scaler, and column names
with open(f'model/columns.pkl', 'rb') as f:
    cols = pickle.load(f)

with open(f'model/scaler.pkl', 'rb') as f:
    sc = pickle.load(f)

with open(f'model/model.pkl', 'rb') as f:
    model = pickle.load(f)

app = FastAPI()

@app.post("/predict/")
async def predict(request: Request):
    try:
        data = await request.json()

        df = pd.read_json(data, orient='records')

        preprocessor = Preprocess(
            dataframe=df,
            missing_value_per=0,
            variance_threshold=0,
            min_null_per=0
        )
        test = preprocessor._mapping(df)

        m = MissingValue(test)
        test = m.fill_dataframe()

        x_test = test[cols]
        x_test = sc.transform(x_test)

        optimized_model_path = 'model/optimized_random_forest_model.onnx'

        session = ort.InferenceSession(optimized_model_path)

        # Run inference
        input_name = session.get_inputs()[0].name
        output_name = session.get_outputs()[0].name

        x_test = np.array(x_test, dtype=np.float32)
        # Make predictions
        y_pred = session.run([output_name], {input_name: x_test})
        y_pred = y_pred[0]

        # y_pred = model.predict(x_test)

        return {"predictions": y_pred.tolist()}
        # return {"predictions": y_pred.tolist()}

    except Exception as e:
        return {"error": str(e)}

# Run the application using: uvicorn main:app --reload