# AUTOGENERATED! DO NOT EDIT! File to edit: 03_tab.stats.ipynb (unless otherwise specified).

__all__ = ['Categorify', 'setups', 'FillMissing']

# Cell
from fastai.tabular.all import *

# Cell
from fastai.tabular.core import _apply_cats

# Cell
class Categorify(TabularProc):
    "Transform the categorical variables to something similar to `pd.Categorical`"
    order = 1
    def __init__(self, classes=None):
        if classes is None: classes = defaultdict(L)
        store_attr()
        super().__init__()
    def setups(self, to):
        for n in to.cat_names:
            if n not in self.classes or is_categorical_dtype(to[n]):
                self.classes[n] = CategoryMap(to.iloc[:,n].items, add_na=n)

    def encodes(self, to): to.transform(to.cat_names, partial(_apply_cats, self.classes, 1))
    def decodes(self, to): to.transform(to.cat_names, partial(_decode_cats, self.classes))
    def __getitem__(self,k): return self.classes[k]

# Cell
@patch
def __init__(x:Normalize, mean=None, std=None, axes=(0,2,3), means=None, stds=None): store_attr()

# Cell
@Normalize
def setups(self, to:Tabular):
    if self.means is None:
        self.means = dict(getattr(to, 'train', to).conts.mean())
        self.stds  = dict(getattr(to, 'train', to).conts.std(ddof=0)+1e-7)
    else:
        for name in getattr(to, 'train', to).conts.columns:
            if name not in self.means.keys():
                self.means[name] = getattr(to, 'train', to)[name].mean()
            if name not in self.stds.keys():
                self.stds[name] = getattr(to, 'train', to)[name].std(ddof=0)+1e-7

    return self(to)

# Cell
class FillMissing(TabularProc):
    "Fill the missing values in continuous columns."
    def __init__(self, fill_strategy=FillStrategy.median, add_col=True, fill_vals=None, na_dict=None):
        if fill_vals is None: fill_vals = defaultdict(int)
        store_attr()

    def setups(self, to):
        missing = pd.isnull(to.conts).any()
        if self.na_dict is None:
            store_attr(na_dict={n:self.fill_strategy(to[n], self.fill_vals[n])
                                for n in missing[missing].keys()})
        else:
            diff_cols = list(set(to.cont_names) - set(self.na_dict.keys()))
            for key in diff_cols:
                self.na_dict[key] = self.fill_strategy(to[key], self.fill_vals[key])
        self.fill_strategy = self.fill_strategy.__name__

    def encodes(self, to):
        missing = pd.isnull(to.conts)
        for n in missing.any()[missing.any()].keys():
            assert n in self.na_dict, f"nan values in `{n}` but not in setup training set"
        for n in self.na_dict.keys():
            to[n].fillna(self.na_dict[n], inplace=True)
            if self.add_col:
                to.loc[:,n+'_na'] = missing[n]
                if n+'_na' not in to.cat_names: to.cat_names.append(n+'_na')