#! /usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Sun Jun 17 01:19:29 2018
sklearn supplemental classes

あると便利な sklearn クラス
@author: ks
"""
import numpy as np
import warnings
from sklearn.base import BaseEstimator, RegressorMixin, TransformerMixin, clone
from sklearn.utils.metaestimators import _BaseComposition
from sklearn.utils.validation import check_is_fitted
from sklearn.utils import check_array


class ChainTransformedTargetRegressor(BaseEstimator, RegressorMixin):
    """
    transformers をリストで与える以外は TransformedTargetRegressor と同じ
    """
    def __init__(self, regressor=None, transformers=None):
        self.regressor = regressor
        self.transformers = transformers

    def _fit_transformers(self, y):
        self.transformers_ = self.transformers
        for i in range(len(self.transformers)):
            y = self.transformers_[i].fit_transform(y)
        # XXX: sample_weight is not currently passed to the
        # transformer. However, if transformer starts using sample_weight, the
        # code should be modified accordingly. At the time to consider the
        # sample_prop feature, it is also a good use case to be considered.

    def fit(self, X, y, sample_weight=None):
        """Fit the model according to the given training data.
        Parameters
        ----------
        X : {array-like, sparse matrix}, shape (n_samples, n_features)
            Training vector, where n_samples is the number of samples and
            n_features is the number of features.
        y : array-like, shape (n_samples,)
            Target values.
        sample_weight : array-like, shape (n_samples,) optional
            Array of weights that are assigned to individual samples.
            If not provided, then each sample is given unit weight.
        Returns
        -------
        self : object
        """
        y = check_array(y, accept_sparse=False, force_all_finite=True,
                        ensure_2d=False, dtype='numeric')

        # store the number of dimension of the target to predict an array of
        # similar shape at predict
        self._training_dim = y.ndim

        # transformers are designed to modify X which is 2d dimensional, we
        # need to modify y accordingly.
        if y.ndim == 1:
            y_2d = y.reshape(-1, 1)
        else:
            y_2d = y
        self._fit_transformers(y_2d)

        if self.regressor is None:
            from sklearn.linear_model import LinearRegression
            self.regressor_ = LinearRegression()
        else:
            self.regressor_ = clone(self.regressor)

        # transform y and convert back to 1d array if needed
        y_trans = y_2d
        for i in range(len(self.transformers_)):
            y_trans = self.transformers_[i].transform(y_trans)
        # FIXME: a FunctionTransformer can return a 1D array even when validate
        # is set to True. Therefore, we need to check the number of dimension
        # first.
        if y_trans.ndim == 2 and y_trans.shape[1] == 1:
            y_trans = y_trans.squeeze(axis=1)
        if sample_weight is None:
            self.regressor_.fit(X, y_trans)
        else:
            self.regressor_.fit(X, y_trans, sample_weight=sample_weight)

        return self

    def predict(self, X):
        """Predict using the base regressor, applying inverse.
        The regressor is used to predict and the ``inverse_func`` or
        ``inverse_transform`` is applied before returning the prediction.
        Parameters
        ----------
        X : {array-like, sparse matrix}, shape = (n_samples, n_features)
            Samples.
        Returns
        -------
        y_hat : array, shape = (n_samples,)
            Predicted values.
        """
        check_is_fitted(self, "regressor_")
        pred = self.regressor_.predict(X)
        if pred.ndim == 1:
            pred_trans = pred.reshape(-1, 1)
        else:
            pred_trans = pred
        for i in reversed(range(len(self.transformers_))):
            pred_trans = self.transformers_[i].inverse_transform(pred_trans)
        if (self._training_dim == 1 and
                pred_trans.ndim == 2 and pred_trans.shape[1] == 1):
            pred_trans = pred_trans.squeeze(axis=1)

        return pred_trans


class ColumnTransformer(_BaseComposition, TransformerMixin):
    """"
    sklearn 0.20dev の同名クラスが現時点で動かないので代用
    sklearn で実装されたら消す
    --------
    transformers: タプルリスト。タプルは (transformer 名, transformer, 対象列) で表す
    remainder: 一度も指定されなかった列をどうするか。 'patssrough' でそのまま出力。
              'drop' で出力しない。デフォルト: 'passthrough'
    """
    def __init__(self, transformers, remainder='passthrough'):
        self.transformers = transformers
        self.remainder = remainder

    def fit(self, X, y=None):
        X = X.copy()
        if type(X).__name__ == 'DataFrame':
            for n, t, c in self.transformers:
                t.fit(X=X[c].values, y=y)
        elif type(X).__name__ == 'ndarray':
            for n, t, c in self.transformers:
                t.fit(X=X[:, c], y=y)
        return self

    def transform(self, X, y=None):
        if type(X).__name__ == 'DataFrame':
            X_ = [t.transform(X[c].values, y) for n, t, c in self.transformers]
        elif type(X).__name__ == 'ndarray':
            X_ = [t.transform(X[:, c], y) for n, t, c in self.transformers]
        if self.remainder == 'passthrough':
            remainders = tuple(set(range(X.shape[1])) - set(
                    [c for _, _, cols in self.transformers for c in cols]))
            if remainders is not None:
                X_ += [X[:, remainders]]
        elif self.remainder == 'drop':
            pass
        else:
            raise ValueError('paththrough must be `passthrough` or `drop`.')
        print([x.shape for x in X_])
        X_ = np.hstack(X_)
        return X_

    def fit_transform(self, X, y=None):
        return self.fit(X, y).transform(X, y)


class LabelEncoderM(BaseEstimator, TransformerMixin):
    """
    sklearn.preprocessing.LabelEncoder の複数列バージョン
    ------
    col_index: 変換する列を指定。 None ならば全列が対象。デフォルト: None
    handle_unknown: transform 時に、fit の入力 X にないラベルがあった場合の対応。
                    'impute' で補完、 'raise' でエラーを, 'warn' で警告を出して補完
                    デフォルト: warn
    impute_value: transform 時の補完値。handle_unknown で
                  'warn', 'impute' のいずれかを指定したときに有効。default: 0
    """
    def __init__(self, col_index=None, handle_unknown='warn', impute_value=0):
        self.col_index = col_index
        self.handle_unknown = handle_unknown
        self.impute_value = impute_value

    def fit(self, X, y=None):
        X = X.copy()
        if type(X).__name__ == 'DataFrame':
            if self.col_index is not None:
                X = X[self.col_index].values
        elif type(X).__name__ == 'ndarray':
            if self.col_index is not None:
                X = X[:, self.col_index]
        else:
            X = np.array(X)
        self.labels_ = [
                {k: v for v, k in enumerate(np.unique(c))} for c in X.T]
        return self

    def transform(self, X, y=None):
        X = X.copy()
        if type(X).__name__ == 'DataFrame':
            if self.col_index is not None:
                X = X[self.col_index].values
        elif type(X).__name__ == 'ndarray':
            if self.col_index is not None:
                X = X[:, self.col_index]
        else:
            X = np.array(X)
        X = X.astype('object')
        if len(self.labels_) != X.shape[1]:
            raise IndexError('dimension of classes and X is wrong')
        for i in range(X.shape[1]):
            X[:, i] = [self.labels_[i].get(k) for k in X[:, i]]
            if self.handle_unknown in ['raise', 'warn'] and np.isin(
                    None, X[:, i]):
                if self.handle_unknown == 'raise':
                    raise ValueError("unknown label found at "
                                     "%s' th feature column" % i)
                else:
                    warnings.warn("unknown label found at "
                                  "%s' th feature column" % i)
            X[np.equal(X[:, i], None), i] = self.impute_value
        return X.astype(float)

    def fit_transform(self, X, y=None):
        return self.fit(X, y).transform(X, y)