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, 'sklearn.utils.multiclass.check_classification_targets', 'check_classification_targets', (['y'], {}), False, 'from sklearn.utils.multiclass import check_classification_targets, type_of_target\n'), (47, 'numpy.unique', 'np.unique', (['X'], {}), True, 'import numpy as np\n'), (73, 'numpy.unique', 'np.unique', (['X'], {}), True, 'import numpy as np\n'), (106, 'numpy.array', 'np.array', (['votes_per_example'], {}), True, 'import numpy as np\n'), (142, 'sklearn.utils.validation.check_X_y', 'check_X_y', (['X', 'y'], {}), False, 'from sklearn.utils.validation import check_X_y, check_array, NotFittedError\n'), (143, 'sklearn.neighbors.KNeighborsClassifier', 'KNeighborsClassifier', (['self.k'], {}), False, 'from sklearn.neighbors import KNeighborsClassifier\n'), (160, 'sklearn.utils.validation.check_array', 'check_array', (['X'], {}), False, 'from sklearn.utils.validation import check_X_y, check_array, NotFittedError\n'), (190, 'sklearn.utils.validation.check_X_y', 'check_X_y', (['X', 'y'], {}), False, 'from sklearn.utils.validation import check_X_y, check_array, NotFittedError\n'), (192, 'tensorflow.keras.Sequential', 'keras.Sequential', ([], {}), False, 'from tensorflow import keras\n'), (228, 'sklearn.utils.validation.check_array', 'check_array', (['X'], {}), False, 'from sklearn.utils.validation import check_X_y, check_array, NotFittedError\n'), (255, 'numpy.mean', 'np.mean', (['y'], {}), True, 'import numpy as np\n'), (257, 'numpy.unique', 'np.unique', (['X'], {}), True, 'import numpy as np\n'), (285, 'numpy.array', 'np.array', (['votes_per_example'], {}), True, 'import numpy as np\n'), (37, 'sklearn.utils.multiclass.type_of_target', 'type_of_target', (['y'], {}), False, 'from sklearn.utils.multiclass import check_classification_targets, type_of_target\n'), (42, 'numpy.unique', 'np.unique', (['y'], {}), True, 'import numpy as np\n'), (43, 'numpy.ones', 'np.ones', (['num_classes'], {}), True, 'import numpy as np\n'), (121, 'numpy.where', 'np.where', (['(posteriors == 0)'], {}), True, 'import numpy as np\n'), (194, 'keras.layers.Dense', 'layers.Dense', (['(1)'], {'activation': '"""linear"""', 'input_shape': '(X.shape[1],)', 'name': '"""transform_to_vote"""'}), False, 'from keras import layers\n'), (48, 'numpy.where', 'np.where', (['(X == leaf_id)'], {}), True, 'import numpy as np\n'), (74, 'numpy.where', 'np.where', (['(X == leaf_id)'], {}), True, 'import numpy as np\n'), (202, 'tensorflow.keras.optimizers.Adam', 'keras.optimizers.Adam', (['self.lr'], {}), False, 'from tensorflow import keras\n'), (258, 'numpy.where', 'np.where', (['(X == leaf_id)'], {}), True, 'import numpy as np\n'), (260, 'numpy.mean', 'np.mean', (['y[idxs_in_leaf]'], {}), True, 'import numpy as np\n'), (50, 'numpy.unique', 'np.unique', (['y'], {}), True, 'import numpy as np\n'), (52, 'numpy.array', 'np.array', (['class_counts'], {}), True, 'import numpy as np\n'), (52, 'numpy.sum', 'np.sum', (['class_counts'], {}), True, 'import numpy as np\n'), (78, 'numpy.array', 'np.array', (['label_counts'], {}), True, 'import numpy as np\n'), (78, 'numpy.sum', 'np.sum', (['label_counts'], {}), True, 'import numpy as np\n'), (208, 'tensorflow.keras.callbacks.EarlyStopping', 'keras.callbacks.EarlyStopping', ([], {'patience': '(20)', 'monitor': '"""val_loss"""'}), False, 'from tensorflow import keras\n'), (50, 'numpy.where', 'np.where', (['(y[idxs_in_leaf] == y_val)'], {}), True, 'import numpy as np\n'), (56, 'numpy.unique', 'np.unique', (['y'], {}), True, 'import numpy as np\n'), (76, 'numpy.where', 'np.where', (['(y[idxs_in_leaf, j] == 1)'], {}), True, 'import numpy as np\n')]