"""
Jupyter Helpers or Kaggle Credit Risk Data

%matplotlib inline
%load_ext autoreload
%autoreload 2
"""


import pandas as pd
import numpy as np
import datetime as dt
import matplotlib.pyplot as plt
import os
import glob
import shutil
import json
import argparse

from sklearn.preprocessing import OneHotEncoder, LabelEncoder
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler


def load_train_data_train_test_split(train_fn, balance_train_classes=False, test_size=0.33, random_state=1234):
    """
    Do a simple train/test split + standard scaling

    :param train_fn:
    :param balance_train_classes:
    :return:
    """
    print("Loading", train_fn)
    train_df = pd.read_csv(train_fn)

    skipcols = ("SK_ID_CURR", "TARGET")
    traincols = [c for c in train_df.columns if c not in skipcols]
    feature_cols = [c for c in train_df.columns if c not in skipcols]

    print(train_df['TARGET'].value_counts())

    X = train_df[feature_cols].values
    y = train_df['TARGET'].values

    scaler = StandardScaler()
    X = scaler.fit_transform(X)

    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=test_size, random_state=random_state)

    if balance_train_classes:
        X_pos = X_train[y_train==1]
        X_neg = X_train[y_train==0]
        num_pos = X_pos.shape[0]

        print("Num pos samples in train set (before)", num_pos)

        X_neg = X_neg[:num_pos, :]

        y_neg = np.zeros(X_neg.shape[0])
        y_pos = np.ones(num_pos)

        X_train = np.concatenate((X_neg, X_pos))
        y_train = np.concatenate((y_neg, y_pos))

        rnd = np.arange(X_train.shape[0])
        np.random.shuffle(rnd)
        print("rnd.shape", rnd.shape)
        X_train = X_train[rnd, :]
        y_train = y_train[rnd]

    print("X_train.shape", X_train.shape)
    print("y_train.shape", y_train.shape)
    print("y_train labels counts", np.unique(y_train, return_counts=True))

    print("X_test.shape", X_test.shape)
    print("y_test.shape", y_test.shape)
    print("y_test labels counts", np.unique(y_test, return_counts=True))

    print("feature_cols.shape", len(feature_cols))

    return (X_train, X_test, y_train, y_test, scaler)


def load_train_and_kaggle_submission(test_fn, scaler):
    """
    Load the Kaggle submission CSV JOINED with the the features.
    Scale the features using the scaler created by load_train_data_train_test_split()

    :param test_fn: File path to kaggle submission CSV
    :param scaler: StandardScaler object (already fit())
    :return:
    """
    print("[load_train_and_kaggle_submission] - start")
    print("Loading", test_fn)
    test_df = pd.read_csv(test_fn)

    skipcols = ("SK_ID_CURR", "TARGET")
    feature_cols = [c for c in test_df.columns if c not in skipcols]
    X_test = test_df[feature_cols].values

    print("scale(test)")
    X_test = scaler.transform(X_test)

    # test submission data
    print("X_test.shape", X_test.shape)

    # create kaggle submission dataframe, with only the Primary Key
    submission_df = test_df[["SK_ID_CURR"]].copy()
    submission_df["TARGET"] = None
    print("[load_train_and_kaggle_submission] - finished")

    return (X_test, submission_df)


def show_null_value_cols(df):
    """
    Prints columns that contain NULLs

    :param df: Pandas dataframe
    :return: list of column names with NULLs
    """
    # clean nulls
    col_has_nulls = df.isna().any()
    null_col_names = df.columns[col_has_nulls].tolist()
    null_col_dtypes = df.dtypes[col_has_nulls].tolist()
    null_cols = [r for r in zip(null_col_names, null_col_dtypes)]

    # list the null cols
    for colname, coltype in null_cols:
        print("*" * 10)
        if coltype == "float64":
            print(colname, "float")
            print("\t mean()", df[colname].mean())
        elif coltype == 'object':
            vals = list(df[colname].unique())
            print(colname, "object")
            print("\t", vals)

        print("")

    return null_cols


def fix_null_values(df, null_cols, impute="mean"):
    """
    Replace NULL values for columns using method determined by `impute`

    :param df: pandas dataframe
    :param null_cols: null column list generated by show_null_value_cols()
    :param impute: `mean` : replace with column mean. For `object` columns, replace with 'UNKNOWN'
    :return:
    """
    # list the null cols
    for colname, coltype in null_cols:
        # print("*" * 10)
        if coltype in ("float64", "int64"):
            print("fix_null_values", colname, "float")
            if impute == 'mean':
                val = df[colname].mean()
            else:
                val = 0.0
            df[colname].fillna(val, inplace=True)
        elif coltype == 'object':
            print("fix_null_values", colname, "object")
            df[colname].fillna("UNKNOWN", inplace=True)
        else:
            raise RuntimeError("Unknown type: {}".format(coltype))

        print("")

    _ = show_null_value_cols(df=df)


def fix_number_null_values(df, null_cols, impute="mean"):
    """
    Replace NULL values for columns using method determined by `impute`. Skip `object` columns.

    :param df: pandas dataframe
    :param null_cols: null column list generated by show_null_value_cols()
    :param impute: `mean` : replace with column mean. For `object` columns, skip
    :return:
    """

    # list the null cols
    for colname, coltype in null_cols:
        # print("*" * 10)
        if coltype in ("float64", "int64"):
            print("has null", colname, "float")
            if impute == 'mean':
                val = df[colname].mean()
            else:
                val = 0.0
            df[colname].fillna(val, inplace=True)
        elif coltype == 'object':
            continue
        else:
            raise RuntimeError("Unknown type: {}".format(coltype))

        print("")

    null_cols = show_null_value_cols(df=df)


def do_data_cleaning(input_csv, output_csv, impute="mean"):
    """
    Find columns with NULL values, then `impute` those columns

    :param input_csv: File path to raw data
    :param output_csv: File path to write `clean` data
    :param impute: method to use to fix NULL values
    :return:
    """
    print("Reading", input_csv)
    df = pd.read_csv(input_csv)

    print("Find null columns")
    null_cols = show_null_value_cols(df=df)
    fix_null_values(df=df, null_cols=null_cols, impute=impute)

    print("Sanity check")
    null_cols = show_null_value_cols(df=df)
    assert len(null_cols) == 0
    print("Sanity check passed")

    print("Write to", output_csv)
    df.to_csv(output_csv, index=False)
    print("do_data_cleaning - done")


def generate_encoders(df, object_cols=None):
    """
    Generate onehot, label encoders from unique values in the columns. Onehot if > 2 values, else LabelEncode

    :param df: Pandas dataframe
    :param object_cols: columns of type 'object' that we want to encode
    :return: Tuple of dictionaries mapping column names to encoders
    """

    one_hot_encoders_di = {}
    label_encoders_di = {}
    if object_cols is None:
        object_cols = df.select_dtypes('object').apply(pd.Series.nunique, axis=0).to_dict()

    for colname, num_unique in object_cols.items():
        vals = list(df[colname].unique())

        if num_unique > 2:
            #         print ("onehot:",colname,":",vals)
            vals_ = np.array(vals).reshape(-1, 1)
            enc = OneHotEncoder(handle_unknown='ignore', sparse=False)
            enc.fit(vals_)
            one_hot_encoders_di[colname] = {'encoder': enc,
                                            'colnames': ["onehot__{}_{}".format(colname, i) for i in range(num_unique)]}
            print("onehot:", colname, ":", enc.categories_)
        else:
            enc = LabelEncoder()
            enc.fit(vals)
            label_encoders_di[colname] = {'encoder': enc,
                                          'colnames': ["label__{}".format(colname)]}
            print("label:", colname, ":", enc.classes_)

    return one_hot_encoders_di, label_encoders_di


def add_onehot_col(df, one_hot_encoders_di, output_feat_dir, idxcol="SK_ID_CURR", drop=True, filename_prefix="",
                   force=False, write_col_csv=False):
    """
    Add onehot col to the input dataframe

    :param df: Pandas dataframe
    :param one_hot_encoders_di: dictionary mapping column names to the encoders
    :param output_feat_dir: output directory if we want to store CSV files containing only the new column
    :param idxcol: Primary Key col from df
    :param drop: whether we want to drop idxcol from the return df
    :param filename_prefix: string used to disambiguate columns with same name from different datasets.
    :param force: If true, overwrite existing data
    :param write_col_csv: True, if we want to save the onehot cols in a CSV
    :return: Pandas dataframe
    """
    df.set_index(idxcol, drop=False, inplace=True)
    onehot_df = df

    for colname in one_hot_encoders_di:
        output_csv_gz = os.path.join(output_feat_dir, "{}onehot__colname__{}.csv.gz".format(filename_prefix,
                                                                                            colname))
        if os.path.isfile(output_csv_gz) and not force:
            print("[add_onehot_col] - skipping: {}".format(colname))
            continue

        print("[add_onehot_col] - adding: {}".format(colname))
        enc = one_hot_encoders_di[colname]["encoder"]
        colnames = one_hot_encoders_di[colname]["colnames"]
        values = df[colname].values

        print(colname, "values.shape", values.shape)
        onehot_values = enc.transform(values.reshape(-1, 1))

        print(colname, "onehot_values.shape", onehot_values.shape)

        # write the onehot values to disk
        temp_df = pd.DataFrame(onehot_values, columns=colnames,
                                 index=df.index)
        temp_df.index.name = idxcol
        temp_df[colname] = values

        if write_col_csv:
            temp_df.to_csv(output_csv_gz, compression="gzip", index=False)
            #         column_names = onehot_df.columns.tolist()
            #         print(json.dumps(sorted(column_names), indent=2))
            print("Wrote to", output_csv_gz)
            print("")

        # print("drop", colname, type(colname))
        temp_df = temp_df.drop(colname, axis=1)

        onehot_df = pd.concat([onehot_df, temp_df], axis=1, copy=False)
        if drop:
            onehot_df = onehot_df.drop(colname, axis=1)
            #         column_names = onehot_df.columns.tolist()
            #         print(json.dumps(sorted(column_names), indent=2))
            #         break

    return onehot_df


def add_label_col(df, label_encoders_di, output_feat_dir, idxcol="SK_ID_CURR", drop=True, filename_prefix="",
                  force=False, write_col_csv=False):
    """
    Add onehot col to the input dataframe

    :param df: Pandas dataframe
    :param label_encoders_di: dictionary mapping column names to the encoders
    :param output_feat_dir: output directory if we want to store CSV files containing only the new column
    :param idxcol: Primary Key col from df
    :param drop: whether we want to drop idxcol from the return df
    :param filename_prefix: string used to disambiguate columns with same name from different datasets.
    :param force: If true, overwrite existing data
    :param write_col_csv: True, if we want to save the onehot cols in a CSV
    :return: Pandas dataframe
    """
    df.set_index(idxcol, drop=False, inplace=True)
    label_df = df

    for colname in label_encoders_di:
        output_csv_gz = os.path.join(output_feat_dir, "{}label__colname__{}.csv.gz".format(filename_prefix,
                                                                                           colname))
        if os.path.isfile(output_csv_gz) and not force:
            print("[add_label_col] - skipping: {}".format(colname))
            continue

        print("[add_label_col] - adding: {}".format(colname))
        enc = label_encoders_di[colname]["encoder"]
        colnames = label_encoders_di[colname]["colnames"]
        values = df[colname].values

        print(colname, "values.shape", values.shape)
        #         label_values = enc.transform(values.reshape(-1, 1))
        label_values = enc.transform(values)

        print(colname, "onehot_values.shape", label_values.shape)

        # write the onehot values to disk
        temp_df = pd.DataFrame(label_values, columns=colnames,
                                index=df.index)
        temp_df.index.name = idxcol
        temp_df[colname] = values

        if write_col_csv:
            temp_df.to_csv(output_csv_gz, compression="gzip", index=False)
            print("Wrote to", output_csv_gz)
            print("")

        # print("drop", colname, type(colname))
        temp_df = temp_df.drop(colname, axis=1)
        label_df = pd.concat([label_df, temp_df], axis=1, copy=False)
        if drop:
            label_df = label_df.drop(colname, axis=1)

    return label_df


def rename_cols_with_feat_name_prefix(feat_code, colnames, df, idxcol="SK_ID_CURR"):
    """
    Rename columns with the `feat_code` prefix to disambiguate columns with same name from
    different datasets.

    :param feat_code: string to prefix column names
    :param colnames: columns to change
    :param df: pandas dataframe
    :param idxcol: Primary key column
    :return:
    """
    df.set_index(idxcol, drop=True, inplace=True)

    # FEAT_CODE = "CCB"
    rename_cols = {}
    for colname in colnames:
        rename_cols[colname] = "{}_{}".format(feat_code, colname)

    df.rename(index=int, inplace=True, columns=rename_cols)

    df.reset_index(inplace=True)

    return list(rename_cols.values())
