from typing import Dict, List, Union
import os
import joblib
import numpy as np

features = ["fixed_acidity", "volatile_acidity", "citric_acid", "residual_sugar", "chlorides", "free_sulfur_dioxide", "total_sulfur_dioxide", "density", "pH", "sulphates", "alcohol"]

class EndpointHandler():
    def __init__(self, path=""):
        self.model = joblib.load("sklearn_model.joblib")

    def __call__(
        self, inputs: Dict[str, Dict[str, List[Union[str, float]]]]
    ) -> List[Union[str, float]]:
        """
        Args:
            inputs (:obj:`dict`):
                a dictionary containing a key 'data' mapping to a dict in which
                the values represent each column.
        Return:
            A :obj:`list` of floats or strings: The classification output for each row.
        """
        data = inputs["data"]
        X = np.array([data[i] for i in features])
        X = np.transpose(X, (1, 0))
        return self.model.predict(X)