# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/data_utils/preprocessing.ipynb.

# %% ../../nbs/data_utils/preprocessing.ipynb 3
from __future__ import annotations
import numpy as np
import einops

# %% auto 0
__all__ = ['DataPreprocessor', 'MinMaxScaler', 'EncoderPreprocessor', 'OrdinalPreprocessor', 'OneHotEncoder']

# %% ../../nbs/data_utils/preprocessing.ipynb 5
def _check_xs(xs: np.ndarray, name: str):
    if xs.ndim > 2 or (xs.ndim == 2 and xs.shape[1] != 1):
        raise ValueError(f"`{name}` only supports array with a single feature, but got shape={xs.shape}.")

# %% ../../nbs/data_utils/preprocessing.ipynb 6
class DataPreprocessor:

    def __init__(
        self, 
        name: str = None # The name of the preprocessor. If None, the class name will be used.
    ):
        """Base class for data preprocessors."""
        self.name = name or self.__class__.__name__
    
    def fit(self, xs, y=None):
        """Fit the preprocessor with `xs` and `y`."""
        raise NotImplementedError
    
    def transform(self, xs):
        """Transform `xs`."""
        raise NotImplementedError
    
    def fit_transform(self, xs, y=None):
        """Fit the preprocessor with `xs` and `y`, then transform `xs`."""
        self.fit(xs, y)
        return self.transform(xs)
    
    def inverse_transform(self, xs):
        """Inverse transform `xs`."""
        raise NotImplementedError
    
    def to_dict(self) -> dict:
        """Convert the preprocessor to a dictionary."""
        raise NotImplementedError
    
    def from_dict(self, params: dict):
        """Load the attributes of the preprocessor from a dictionary."""
        raise NotImplementedError
        
    __ALL__ = ["fit", "transform", "fit_transform", "inverse_transform", "to_dict", "from_dict"]

# %% ../../nbs/data_utils/preprocessing.ipynb 7
class MinMaxScaler(DataPreprocessor): 
    def __init__(self):
        super().__init__(name="minmax")
        
    def fit(self, xs, y=None):
        _check_xs(xs, name="MinMaxScaler")
        self.min_ = xs.min(axis=0)
        self.max_ = xs.max(axis=0)
        return self
    
    def transform(self, xs):
        return (xs - self.min_) / (self.max_ - self.min_)
    
    def inverse_transform(self, xs):
        return xs * (self.max_ - self.min_) + self.min_
    
    def from_dict(self, params: dict):
        self.min_ = params["min_"]
        self.max_ = params["max_"]
        return self
    
    def to_dict(self) -> dict:
        return {"min_": self.min_, "max_": self.max_}

# %% ../../nbs/data_utils/preprocessing.ipynb 13
def _unique(xs):
    if xs.dtype == object:
        # Note: np.unique does not work with object dtype
        # We will enforce xs to be string type
        # It assumes that xs is a list of strings, and might not work
        # for other cases (e.g., list of string and numbers)
        return np.unique(xs.astype(str))
    return np.unique(xs)

# %% ../../nbs/data_utils/preprocessing.ipynb 14
class EncoderPreprocessor(DataPreprocessor):
    """Encode categorical features as an integer array."""
    def _fit(self, xs, y=None):
        _check_xs(xs, name="EncoderPreprocessor")
        self.categories_ = _unique(xs)

    def _transform(self, xs):
        """Transform data to ordinal encoding."""
        if xs.dtype == object:
            xs = xs.astype(str)
        ordinal = np.searchsorted(self.categories_, xs)
        # return einops.rearrange(ordinal, 'k n -> n k')
        return ordinal
    
    def _inverse_transform(self, xs):
        """Transform ordinal encoded data back to original data."""
        return self.categories_[xs.T].T
    
    def from_dict(self, params: dict):
        self.categories_ = params["categories_"]
        return self
    
    def to_dict(self) -> dict:
        return {"categories_": self.categories_}

# %% ../../nbs/data_utils/preprocessing.ipynb 15
class OrdinalPreprocessor(EncoderPreprocessor):
    """Ordinal encoder for a single feature."""
    
    def fit(self, xs, y=None):
        self._fit(xs, y)
        return self
    
    def transform(self, xs):
        if xs.ndim == 1:
            raise ValueError(f"OrdinalPreprocessor only supports 2D array with a single feature, "
                             f"but got shape={xs.shape}.")
        return self._transform(xs)
    
    def inverse_transform(self, xs):
        return self._inverse_transform(xs)

# %% ../../nbs/data_utils/preprocessing.ipynb 17
class OneHotEncoder(EncoderPreprocessor):
    """One-hot encoder for a single categorical feature."""
    
    def fit(self, xs, y=None):
        self._fit(xs, y)
        return self

    def transform(self, xs):
        if xs.ndim == 1:
            raise ValueError(f"OneHotEncoder only supports 2D array with a single feature, "
                             f"but got shape={xs.shape}.")
        xs_int = self._transform(xs)
        one_hot_feats = np.eye(len(self.categories_))[xs_int]
        return einops.rearrange(one_hot_feats, 'n k d -> n (k d)')

    def inverse_transform(self, xs):
        xs_int = np.argmax(xs, axis=-1)
        return self._inverse_transform(xs_int).reshape(-1, 1)
