from .storage import FeatureConstructor
import pandas as pd
import numpy as np
from ..storage.dataframe import DataFrame as KTDF
from ..zoo.cluster import KMeansFeaturizer
from ..utils import list_hash


def empty_like(df):
    if isinstance(df, KTDF):
        return KTDF(df[[]].copy(), df.train, df.encoders, df.slice_id)
    else:
        return KTDF(df[[]].copy())


identity = FeatureConstructor(lambda df: df, cache_default=False)
identity.stl = True
identity.source = "stl.identity"


def merge(dfs):
    if len(dfs) == 1:
        return dfs[0]
    return pd.concat([df.df if isinstance(df, KTDF) else df for df in dfs], axis=1)


def column_selector(columns):
    def __col_selector(df):
        return df[columns]

    return FeatureConstructor(__col_selector, cache_default=False)


def column_dropper(columns):
    def __col_dropper(df):
        return df.drop(columns, axis=1)

    return FeatureConstructor(__col_dropper, cache_default=False)


def concat(funcs):
    def __concat(df):
        return merge([func(df) for func in funcs])

    fc = FeatureConstructor(__concat, cache_default=False)
    fc.source = f"stl.concat([{', '.join([func.__name__ if not func.stl else func.source for func in funcs])}])"
    fc.__name__ = f"concat_{''.join([func.__name__[:2] if not func.stl else func.__name__ for func in funcs])}"
    fc.stl = True
    return fc


def compose(funcs):
    def __compose(df):
        res = empty_like(df)
        for func in funcs:
            res = merge([res, func(res)])
        return res

    fc = FeatureConstructor(__compose, cache_default=False)
    fc.source = f"stl.compose([{', '.join([func.__name__ if not func.stl else func.source for func in funcs])}])"
    fc.__name__ = f"compose_{''.join([func.__name__[:2] if not func.stl else func.__name__ for func in funcs])}"
    fc.stl = True
    return fc


# TODO: reimplement using sklearn.preprocessing.OHE
def make_ohe(cols, sep='_ohe_'):
    def __make_ohe(df):
        return pd.get_dummies(df.df[cols], prefix_sep=sep, columns=cols)

    return FeatureConstructor(__make_ohe, cache_default=False)


def target_encoding(cols, target_col, aggregation='mean', prefix='me_'):
    def __target_encoding(df):
        res = empty_like(df)
        for col in cols:
            if df.train:
                enc = df.groupby(col)[target_col].agg(aggregation)
                df.encoders['__me_' + col] = enc
            else:
                enc = df.encoders['__me_' + col]
            res[prefix + col] = df[col].map(enc)
        return res

    return FeatureConstructor(__target_encoding, cache_default=False)


def target_encode_list(cols, target_col, aggregation='mean', prefix='me_list_'):
    def __target_encode_list(df):
        res = empty_like(df)
        filler = [np.nan]
        for col in cols:
            if df.train:
                enc = (pd.DataFrame([i if type(i) == list else [i] for i in df[col].tolist()], index=df[target_col])
                       .stack()
                       .reset_index(name=col)
                       [[col, target_col]]).groupby(col).agg(aggregation)
                df.encoders['__me_list_' + col] = enc
            else:
                enc = df.encoders['__me_list_' + col]
            res[prefix + col] = df[col].apply(lambda x: enc.loc[filter(lambda i: i in enc.index, x)].values.flatten() if type(x) == list else filler)
        return res

    return FeatureConstructor(__target_encode_list, cache_default=False)


import multiprocessing
import swifter


def apply(dataframe, function, **kwargs):
    """
    Applies function to dataframe faster.
    If n_threads is in kwargs and is greater than 1, applies by multiprocessing.
    :return: same as df.apply(function)
    """
    if 'n_threads' in kwargs:
        n_threads = kwargs.pop('n_threads')
    else:
        n_threads = 1
    if n_threads == 1:
        return dataframe.swifter.apply(function, **kwargs)

    def _apply_df(args):
        df, func, num, kw = args
        return num, df.swifter.apply(func, **kw)

    pool = multiprocessing.Pool(processes=n_threads)
    result = pool.map(_apply_df, [(d, function, i, kwargs) for i, d in enumerate(np.array_split(dataframe, n_threads))])
    pool.close()
    result = sorted(result, key=lambda x: x[0])
    return pd.concat([i[1] for i in result])


def get_categorical(df):
    if df.train:
        cat_features = [col for col in df.columns if df[col].dtype == object]
        df.encoders['__cat_features'] = cat_features
    else:
        cat_features = df.encoders['__cat_features']
    return cat_features


def get_numeric(df):
    if df.train:
        num_features = [col for col in df.columns if df[col].dtype != object]
        df.encoders['__num_features'] = num_features
    else:
        num_features = df.encoders['__num_features']
    return num_features


def discretize(cols, bins, prefix='disc_'):
    def __discretize(df):
        res = empty_like(df)
        for col in cols:
            if df.train:
                res[prefix + str(bins) + '_' + col], enc = pd.cut(df[col], bins, retbins=True)
                df.encoders[f'__disc_{bins}_{col}'] = enc
            else:
                enc = df.encoders[f'__disc_{bins}_{col}']
                res[prefix + str(bins) + '_' + col] = pd.cut(df[col], enc)
        return res

    return FeatureConstructor(__discretize, cache_default=False)


def discretize_quantile(cols, bins, prefix='disc_q_'):
    def __discretize_quantile(df):
        res = empty_like(df)
        for col in cols:
            if df.train:
                res[prefix + str(bins) + '_' + col], enc = pd.qcut(df[col], bins, retbins=True)
                df.encoders[f'__disc_q_{bins}_{col}'] = enc
            else:
                enc = df.encoders[f'__disc_q_{bins}_{col}']
                res[prefix + str(bins) + '_' + col] = pd.cut(df[col], enc)
        return res

    return FeatureConstructor(__discretize_quantile, cache_default=False)


def kmeans_encoding(cols, n_clusters, target_col=None, target_importance=5.0, prefix='km_', **kwargs):
    def __kmeans_encoding(df):
        res = empty_like(df)
        res_column_name = f'{prefix}{n_clusters}_{list_hash(cols, 5)}_{round(target_importance, 4)}'
        if df.train:
            encoder = KMeansFeaturizer(k=n_clusters, target_scale=target_importance, **kwargs)
            if target_col:
                res[res_column_name] = encoder.fit_transform(df[cols].values, df[target_col].values)
            else:
                res[res_column_name] = encoder.fit_transform(df[cols].values)
            df.encoders[f'__kmeans_feat_{n_clusters}_{list_hash(cols, 5)}_{round(target_importance, 4)}'] = encoder
        else:
            encoder = df.encoders[f'__kmeans_feat_{n_clusters}_{list_hash(cols, 5)}_{round(target_importance, 4)}']
            res[res_column_name] = encoder.transform()
        return res

    return FeatureConstructor(__kmeans_encoding, cache_default=False)