import numpy as np

# from sklearn.ensemble import BaggingClassifier
# from sklearn.tree import DecisionTreeClassifier
from sklearn.neighbors import KNeighborsClassifier

from sklearn.utils.validation import (
    check_X_y,
    check_array,
    NotFittedError,
)

from sklearn.utils.multiclass import check_classification_targets, type_of_target

from .base import BaseVoter

from tensorflow import keras
from keras import layers


class TreeClassificationVoter(BaseVoter):
    def __init__(self, finite_sample_correction=False):
        """
        Doc strings here.
        """

        self.finite_sample_correction = finite_sample_correction
        self._is_fitted = False
        self.multilabel = False

    def fit(self, X, y):
        """
        Doc strings here.
        """
        check_classification_targets(y)

        if type_of_target(y) == 'multilabel-indicator':
            # Fit multilabel binary task.
            self.multilabel = True
            return self.fit_multilabel(X, y)

        num_classes = len(np.unique(y))
        self.uniform_posterior = np.ones(num_classes) / num_classes

        self.leaf_to_posterior = {}

        for leaf_id in np.unique(X):
            idxs_in_leaf = np.where(X == leaf_id)[0]
            class_counts = [
                len(np.where(y[idxs_in_leaf] == y_val)[0]) for y_val in np.unique(y)
            ]
            posteriors = np.nan_to_num(np.array(class_counts) / np.sum(class_counts))

            if self.finite_sample_correction:
                posteriors = self._finite_sample_correction(
                    posteriors, len(idxs_in_leaf), len(np.unique(y))
                )

            self.leaf_to_posterior[leaf_id] = posteriors

        self._is_fitted = True

        return self

    def fit_multilabel(self, X, y):

        num_labels = y.shape[1]
        self.uniform_posterior = y.sum(axis=0) / len(y)
        
        # Each posterior is now a num_labels size vector or binary probabilities.
        self.leaf_to_posterior = {}

        for leaf_id in np.unique(X):
            idxs_in_leaf = np.where(X == leaf_id)[0]
            label_counts = [
                len(np.where(y[idxs_in_leaf, j] == 1)[0]) for j in range(num_labels)
            ]
            posteriors = np.nan_to_num(np.array(label_counts) / np.sum(label_counts))

            # TODO: multilabel finite sample correction.

            self.leaf_to_posterior[leaf_id] = posteriors

        self._is_fitted = True

        return self

    def vote(self, X):
        """
        Doc strings here.
        """

        if not self.is_fitted():
            msg = (
                "This %(name)s instance is not fitted yet. Call 'fit' with "
                "appropriate arguments before using this voter."
            )
            raise NotFittedError(msg % {"name": type(self).__name__})

        votes_per_example = []
        for x in X:
            if x in list(self.leaf_to_posterior.keys()):
                votes_per_example.append(self.leaf_to_posterior[x])
            else:
                votes_per_example.append(self.uniform_posterior)
        return np.array(votes_per_example)

    def is_fitted(self):
        """
        Doc strings here.
        """

        return self._is_fitted

    def _finite_sample_correction(posteriors, num_points_in_partition, num_classes):
        """
        encourage posteriors to approach uniform when there is low data
        """
        correction_constant = 1 / (num_classes * num_points_in_partition)

        zero_posterior_idxs = np.where(posteriors == 0)[0]
        posteriors[zero_posterior_idxs] = correction_constant

        posteriors /= sum(posteriors)

        return posteriors


class KNNClassificationVoter(BaseVoter):
    def __init__(self, k, kwargs={}):
        """
        Doc strings here.
        """
        self._is_fitted = False
        self.k = k
        self.kwargs = kwargs

    def fit(self, X, y):
        """
        Doc strings here.
        """
        X, y = check_X_y(X, y)
        self.knn = KNeighborsClassifier(self.k, **self.kwargs)
        self.knn.fit(X, y)
        self._is_fitted = True

        return self

    def vote(self, X):
        """
        Doc strings here.
        """
        if not self.is_fitted():
            msg = (
                "This %(name)s instance is not fitted yet. Call 'fit' with "
                "appropriate arguments before using this transformer."
            )
            raise NotFittedError(msg % {"name": type(self).__name__})

        X = check_array(X)
        return self.knn.predict_proba(X)

    def is_fitted(self):
        """
        Doc strings here.
        """

        return self._is_fitted


class NeuralRegressionVoter(BaseVoter):
    def __init__(
        self, validation_split=0.25, loss="mse", epochs=100, lr=1e-4, verbose=False,
    ):
        """
        Doc strings here.
        """
        self.validation_split = validation_split
        self.loss = loss
        self.epochs = epochs
        self.lr = lr
        self.verbose = verbose
        self._is_fitted = False

    def fit(self, X, y):
        """
        Doc strings here.
        """

        X, y = check_X_y(X, y)

        self.voter = keras.Sequential()
        self.voter.add(
            layers.Dense(
                1,
                activation="linear",
                input_shape=(X.shape[1],),
                name="transform_to_vote",
            )
        )
        self.voter.compile(
            loss=self.loss, metrics=["mae"], optimizer=keras.optimizers.Adam(self.lr)
        )
        self.voter.fit(
            X,
            y,
            epochs=self.epochs,
            callbacks=[keras.callbacks.EarlyStopping(patience=20, monitor="val_loss")],
            verbose=self.verbose,
            validation_split=self.validation_split,
            shuffle=True,
        )

        self._is_fitted = True
        return self

    def vote(self, X):
        """
        Doc strings here.
        """
        if not self.is_fitted():
            msg = (
                "This %(name)s instance is not fitted yet. Call 'fit' with "
                "appropriate arguments before using this transformer."
            )
            raise NotFittedError(msg % {"name": type(self).__name__})

        X = check_array(X)
        return self.voter.predict(X)

    def is_fitted(self):
        """
        Doc strings here.
        """

        return self._is_fitted


class TreeRegressionVoter(BaseVoter):
    def __init__(self):
        """
        Doc strings here.
        """

        self._is_fitted = False


    def fit(self, X, y):
        """
        Doc strings here.
        """
        
        
        self.leaf_to_yhat = {}
        self.global_yhat = np.mean(y)

        for leaf_id in np.unique(X):
            idxs_in_leaf = np.where(X == leaf_id)[0]
            # class_counts = [len(np.where(y[idxs_in_leaf] == y_val)[0]) for y_val in np.unique(y)]
            self.leaf_to_yhat[leaf_id] = np.nan_to_num(np.mean(y[idxs_in_leaf]))

        self._is_fitted = True

        return self


    def vote(self, X):
        """
        Doc strings here.
        """

        if not self.is_fitted():
            msg = (
                    "This %(name)s instance is not fitted yet. Call 'fit' with "
                    "appropriate arguments before using this voter."
            )
            raise NotFittedError(msg % {"name": type(self).__name__})
        
        votes_per_example = []
        for x in X:
            if x in list(self.leaf_to_yhat.keys()):
                votes_per_example.append(self.leaf_to_yhat[x])
            else:
                votes_per_example.append(self.global_yhat)
        return np.array(votes_per_example)


    def is_fitted(self):
        """
        Doc strings here.
        """

        return self._is_fitted
[((35, 8, 35, 39), 'sklearn.utils.multiclass.check_classification_targets', 'check_classification_targets', ({(35, 37, 35, 38): 'y'}, {}), '(y)', False, 'from sklearn.utils.multiclass import check_classification_targets, type_of_target\n'), ((47, 23, 47, 35), 'numpy.unique', 'np.unique', ({(47, 33, 47, 34): 'X'}, {}), '(X)', True, 'import numpy as np\n'), ((73, 23, 73, 35), 'numpy.unique', 'np.unique', ({(73, 33, 73, 34): 'X'}, {}), '(X)', True, 'import numpy as np\n'), ((106, 15, 106, 42), 'numpy.array', 'np.array', ({(106, 24, 106, 41): 'votes_per_example'}, {}), '(votes_per_example)', True, 'import numpy as np\n'), ((142, 15, 142, 30), 'sklearn.utils.validation.check_X_y', 'check_X_y', ({(142, 25, 142, 26): 'X', (142, 28, 142, 29): 'y'}, {}), '(X, y)', False, 'from sklearn.utils.validation import check_X_y, check_array, NotFittedError\n'), ((143, 19, 143, 62), 'sklearn.neighbors.KNeighborsClassifier', 'KNeighborsClassifier', ({(143, 40, 143, 46): 'self.k'}, {}), '(self.k, **self.kwargs)', False, 'from sklearn.neighbors import KNeighborsClassifier\n'), ((160, 12, 160, 26), 'sklearn.utils.validation.check_array', 'check_array', ({(160, 24, 160, 25): 'X'}, {}), '(X)', False, 'from sklearn.utils.validation import check_X_y, check_array, NotFittedError\n'), ((190, 15, 190, 30), 'sklearn.utils.validation.check_X_y', 'check_X_y', ({(190, 25, 190, 26): 'X', (190, 28, 190, 29): 'y'}, {}), '(X, y)', False, 'from sklearn.utils.validation import check_X_y, check_array, NotFittedError\n'), ((192, 21, 192, 39), 'tensorflow.keras.Sequential', 'keras.Sequential', ({}, {}), '()', False, 'from tensorflow import keras\n'), ((228, 12, 228, 26), 'sklearn.utils.validation.check_array', 'check_array', ({(228, 24, 228, 25): 'X'}, {}), '(X)', False, 'from sklearn.utils.validation import check_X_y, check_array, NotFittedError\n'), ((255, 27, 255, 37), 'numpy.mean', 'np.mean', ({(255, 35, 255, 36): 'y'}, {}), '(y)', True, 'import numpy as np\n'), ((257, 23, 257, 35), 'numpy.unique', 'np.unique', ({(257, 33, 257, 34): 'X'}, {}), '(X)', True, 'import numpy as np\n'), ((285, 15, 285, 42), 'numpy.array', 'np.array', ({(285, 24, 285, 41): 'votes_per_example'}, {}), '(votes_per_example)', True, 'import numpy as np\n'), ((37, 11, 37, 28), 'sklearn.utils.multiclass.type_of_target', 'type_of_target', ({(37, 26, 37, 27): 'y'}, {}), '(y)', False, 'from sklearn.utils.multiclass import check_classification_targets, type_of_target\n'), ((42, 26, 42, 38), 'numpy.unique', 'np.unique', ({(42, 36, 42, 37): 'y'}, {}), '(y)', True, 'import numpy as np\n'), ((43, 33, 43, 53), 'numpy.ones', 'np.ones', ({(43, 41, 43, 52): 'num_classes'}, {}), '(num_classes)', True, 'import numpy as np\n'), ((121, 30, 121, 55), 'numpy.where', 'np.where', ({(121, 39, 121, 54): '(posteriors == 0)'}, {}), '(posteriors == 0)', True, 'import numpy as np\n'), ((194, 12, 199, 13), 'keras.layers.Dense', 'layers.Dense', (), '', False, 'from keras import layers\n'), ((48, 27, 48, 49), 'numpy.where', 'np.where', ({(48, 36, 48, 48): '(X == leaf_id)'}, {}), '(X == leaf_id)', True, 'import numpy as np\n'), ((74, 27, 74, 49), 'numpy.where', 'np.where', ({(74, 36, 74, 48): '(X == leaf_id)'}, {}), '(X == leaf_id)', True, 'import numpy as np\n'), ((202, 55, 202, 85), 'tensorflow.keras.optimizers.Adam', 'keras.optimizers.Adam', ({(202, 77, 202, 84): 'self.lr'}, {}), '(self.lr)', False, 'from tensorflow import keras\n'), ((258, 27, 258, 49), 'numpy.where', 'np.where', ({(258, 36, 258, 48): '(X == leaf_id)'}, {}), '(X == leaf_id)', True, 'import numpy as np\n'), ((260, 55, 260, 79), 'numpy.mean', 'np.mean', ({(260, 63, 260, 78): 'y[idxs_in_leaf]'}, {}), '(y[idxs_in_leaf])', True, 'import numpy as np\n'), ((50, 72, 50, 84), 'numpy.unique', 'np.unique', ({(50, 82, 50, 83): 'y'}, {}), '(y)', True, 'import numpy as np\n'), ((52, 39, 52, 61), 'numpy.array', 'np.array', ({(52, 48, 52, 60): 'class_counts'}, {}), '(class_counts)', True, 'import numpy as np\n'), ((52, 64, 52, 84), 'numpy.sum', 'np.sum', ({(52, 71, 52, 83): 'class_counts'}, {}), '(class_counts)', True, 'import numpy as np\n'), ((78, 39, 78, 61), 'numpy.array', 'np.array', ({(78, 48, 78, 60): 'label_counts'}, {}), '(label_counts)', True, 'import numpy as np\n'), ((78, 64, 78, 84), 'numpy.sum', 'np.sum', ({(78, 71, 78, 83): 'label_counts'}, {}), '(label_counts)', True, 'import numpy as np\n'), ((208, 23, 208, 85), 'tensorflow.keras.callbacks.EarlyStopping', 'keras.callbacks.EarlyStopping', (), '', False, 'from tensorflow import keras\n'), ((50, 20, 50, 54), 'numpy.where', 'np.where', ({(50, 29, 50, 53): '(y[idxs_in_leaf] == y_val)'}, {}), '(y[idxs_in_leaf] == y_val)', True, 'import numpy as np\n'), ((56, 55, 56, 67), 'numpy.unique', 'np.unique', ({(56, 65, 56, 66): 'y'}, {}), '(y)', True, 'import numpy as np\n'), ((76, 20, 76, 53), 'numpy.where', 'np.where', ({(76, 29, 76, 52): '(y[idxs_in_leaf, j] == 1)'}, {}), '(y[idxs_in_leaf, j] == 1)', True, 'import numpy as np\n')]