from sklearn.model_selection import train_test_split
from sklearn.impute import SimpleImputer
from sklearn.pipeline import make_pipeline
from sklearn.preprocessing import StandardScaler, OneHotEncoder, OrdinalEncoder
from sklearn.compose import ColumnTransformer
from sklearn.compose import make_column_selector
# set the default output to pandas
from sklearn import set_config
set_config(transform_output='pandas')


def preprocessor_maker(df, target, num_imputer='mean', fill_num = 0, ord_features=[], ord_orders=[]
          , ord_imputer='constant', fill_ord='NA', cat_imputer='constant'
          ,fill_cat='NA',  random_state=42, num_scale = True, ord_scale = True, cat_scale = True ):


  y = df[target]
  X = df.drop([target] , axis=1)
  X_train, X_test, y_train, y_test = train_test_split(X, y, random_state=random_state)

  indicator_arr = [0,0,0]

  if (len(X_train.select_dtypes('number').columns)):
    num_selector = make_column_selector(dtype_include='number')

    if num_imputer == 'constant':
      num_imputer1 = SimpleImputer(strategy='constant', fill_value = fill_num)
    else:
      num_imputer1 = SimpleImputer(strategy=num_imputer)

    if num_scale == True:
      scaler = StandardScaler()
      num_pipe = make_pipeline(num_imputer1, scaler)
    else:
      num_pipe = make_pipeline(num_imputer1)


    num_tuple = ('numeric',num_pipe, num_selector)

    indicator_arr[0] = 1



  # ================================================
  if len(ord_features):

    ord_cols = ord_features


    if ord_imputer == 'constant':
      impute_na_ord = SimpleImputer(strategy='constant', fill_value=fill_ord)
    else:
      impute_na_ord = SimpleImputer(strategy=ord_imputer)


    ord_encoder = OrdinalEncoder(categories=ord_orders)

    # # Making a final scaler to scale category #'s
    # scaler_ord = StandardScaler()

    if ord_scale == True:
      scaler_ord = StandardScaler()
      ord_pipe = make_pipeline(impute_na_ord, ord_encoder, scaler_ord)
    else:
      ord_pipe = make_pipeline(impute_na_ord, ord_encoder)

    # ord_pipe = make_pipeline(impute_na_ord, ord_encoder, scaler_ord)



    ord_tuple = ('ordinal', ord_pipe, ord_cols)
    # indicator_arr.append(ord_tuple)
    indicator_arr[1] = 1
  # ==================================================================

  # (New) Select columns with make_column_selector----------------------------------------------
  if (len(X_train.select_dtypes('object').columns)):
    # cat_selector = make_column_selector(dtype_include='object')

    if len(ord_features):
      ohe_cols = X_train.select_dtypes('object').drop(columns=ord_cols).columns


    else:
      ohe_cols = X_train.select_dtypes('object').columns


    if cat_imputer == 'constant':
      freq_imputer = SimpleImputer(strategy='constant', fill_value=fill_cat)
    else:
      freq_imputer = SimpleImputer(strategy=cat_imputer)
      
    ohe = OneHotEncoder(sparse_output=False, handle_unknown='ignore')


    if cat_scale == True:
      scaler_cat = StandardScaler()
      cat_pipe = make_pipeline(freq_imputer, ohe, scaler_cat)
    else:
      cat_pipe = make_pipeline(freq_imputer, ohe)


    cat_tuple = ('categorical',cat_pipe, ohe_cols)
    indicator_arr[2] = 1


  # =======================================================================


  if indicator_arr[0] == 1 and indicator_arr[1] == 1 and indicator_arr[2] == 1:
    all_tuple = [num_tuple, ord_tuple, cat_tuple]

  elif indicator_arr[0] == 1 and indicator_arr[1] == 1 and indicator_arr[2] == 0:
    all_tuple = [num_tuple, ord_tuple]

  elif indicator_arr[0] == 1 and indicator_arr[1] == 0 and indicator_arr[2] == 1:
    all_tuple = [num_tuple, cat_tuple]

  elif indicator_arr[0] == 0 and indicator_arr[1] == 1 and indicator_arr[2] == 1:
    all_tuple = [ord_tuple, cat_tuple]

  elif indicator_arr[0] == 1 and indicator_arr[1] == 0 and indicator_arr[2] == 0:
    all_tuple = [num_tuple]

  elif indicator_arr[0] == 0 and indicator_arr[1] == 0 and indicator_arr[2] == 1:
    all_tuple = [cat_tuple]

  elif indicator_arr[0] == 0 and indicator_arr[1] == 1 and indicator_arr[2] == 0:
    all_tuple = [ord_tuple]

  else:
    print('Error')



  preprocessor = ColumnTransformer(all_tuple, verbose_feature_names_out=False)

  return preprocessor, X_train, y_train, X_test, y_test


