# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/01_data_preprocess.ipynb.

# %% ../nbs/01_data_preprocess.ipynb 4
from __future__ import annotations
from .utils import *
import matplotlib.pyplot as plt
import seaborn as sns
import pandas as pd
import numpy as np
from statsmodels.tools.tools import add_constant
from statsmodels.stats.outliers_influence import variance_inflation_factor
from sklearn.impute import SimpleImputer
from sklearn.preprocessing import MinMaxScaler,StandardScaler

# %% auto 0
__all__ = ['process_missing_values', 'scale_num_cols', 'one_hot_cat', 'preprocessing_general']

# %% ../nbs/01_data_preprocess.ipynb 5
def process_missing_values(X_train:pd.DataFrame, # Training dataframe
                           X_test:pd.DataFrame=None, # Testing dataframe
                           missing_cols:list|str=[], # A column name having missing values, or a list of such columns
                           missing_vals:list|int|float|str=np.NaN, # A placeholder for missing values, or a list of placeholders for all columns in miss_cols
                           strategies:list|str='median', # The imputation strategy from sklearn, or a list of such values. Currently support 'median','mean','most_frequent'
                           **kwargs):
    "Process columns with missing values using Sklearn SimpleInputer"
    if missing_cols==[]:
        return X_train if X_test is None else (X_train,X_test)
    missing_cols = val2list(missing_cols)
    missing_vals = val2list(missing_vals,len(missing_cols))
    strategies = val2list(strategies,len(missing_cols))
    X_train = X_train.copy()
    if X_test is not None: X_test = X_test.copy()
    for i,c in enumerate(missing_cols):
        imp = SimpleImputer(missing_values=missing_vals[i], strategy=strategies[i])
        X_train[c] = imp.fit_transform(X_train[c].values.reshape(-1,1)).flatten()
        if X_test is not None: X_test[c] = imp.transform(X_test[c].values.reshape(-1,1)).flatten()
    return X_train if X_test is None else (X_train,X_test)

# %% ../nbs/01_data_preprocess.ipynb 12
def scale_num_cols(X_train:pd.DataFrame, # Training dataframe
                   X_test:pd.DataFrame=None, # Testing dataframe
                   num_cols:list|str=[], # Name of the numerical column, or a list of such columns
                   scale_methods:list|str='minmax', # Sklearn scaling method ('minmax' or 'standard'), or a list of such methods        
                    **kwargs):
    "Scale numerical columns using Sklearn"
    if num_cols==[]:
        return X_train if X_test is None else (X_train,X_test)
    num_cols = val2list(num_cols)
    scale_methods = val2list(scale_methods,len(num_cols))
    X_train = X_train.copy()
    if X_test is not None: X_test = X_test.copy()
    for i,c in enumerate(num_cols):
        if scale_methods[i]=='minmax':
            imp = MinMaxScaler()
        elif scale_methods[i]=='standard':
            imp = StandardScaler()
        else:
            raise ValueError('Unrecognized scaling method. Accept methods: minmax and standard')
        X_train[c] = imp.fit_transform(X_train[c].values.reshape(-1,1)).flatten()
        if X_test is not None: X_test[c] = imp.transform(X_test[c].values.reshape(-1,1)).flatten()
    return X_train if X_test is None else (X_train,X_test)

# %% ../nbs/01_data_preprocess.ipynb 16
def one_hot_cat(X_train:pd.DataFrame, # Training dataframe
                X_test:pd.DataFrame=None, # Testing dataframe
                cat_cols:list|str=[], # Name of the categorical columns (non-binary), or a list of such columns
                bi_cols:list|str=[], # Name of the binary column, or a list of such columns
                **kwargs):
    "Perform 'get_dummies' on categorical columns"
    if cat_cols==[] and bi_cols==[]:
        return X_train if X_test is None else (X_train,X_test)
    cat_cols = val2list(cat_cols)
    bi_cols = val2list(bi_cols)
    n_train = X_train.shape[0]
    if X_test is not None:
        X_total = pd.concat([X_train,X_test],axis=0)
    else:
        X_total = X_train.copy()
    if len(cat_cols):
        X_total = pd.get_dummies(X_total,columns=cat_cols,drop_first=False)
    if len(bi_cols):
        X_total = pd.get_dummies(X_total,columns=bi_cols,drop_first=True)
    return X_total if X_test is None else (X_total.iloc[:n_train].copy(), X_total.iloc[n_train:].copy())

# %% ../nbs/01_data_preprocess.ipynb 20
def preprocessing_general(X_train:pd.DataFrame, # Training dataframe
                          X_test:pd.DataFrame=None, # Testing dataframe
                          **kwargs, # Keyword arguments for processing missing values, scaling numerical columns and one-hot-encoding categorical columns
                         ):
    """
    The main preprocessing functions, will perform:
    
    - Fill missing values
    
    - Scale numerical columns
    
    - One-hot encode categorical columns
    
    Remember to put in the appropriate keyword arguments for each of the preprocessings mentioned above
    """
    results = process_missing_values(X_train,X_test,**kwargs)
    results = scale_num_cols(results,**kwargs) if X_test is None else scale_num_cols(*results,**kwargs)
    results = one_hot_cat(results,**kwargs) if X_test is None else one_hot_cat(*results,**kwargs)
    return results
