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', False, 'from sklearn.utils.multiclass import check_classification_targets, type_of_target\n'), (47, 'numpy.unique', 'np.unique', True, 'import numpy as np\n'), (73, 'numpy.unique', 'np.unique', True, 'import numpy as np\n'), (106, 'numpy.array', 'np.array', True, 'import numpy as np\n'), (142, 'sklearn.utils.validation.check_X_y', 'check_X_y', False, 'from sklearn.utils.validation import check_X_y, check_array, NotFittedError\n'), (143, 'sklearn.neighbors.KNeighborsClassifier', 'KNeighborsClassifier', False, 'from sklearn.neighbors import KNeighborsClassifier\n'), (160, 'sklearn.utils.validation.check_array', 'check_array', False, 'from sklearn.utils.validation import check_X_y, check_array, NotFittedError\n'), (190, 'sklearn.utils.validation.check_X_y', 'check_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', False, 'from sklearn.utils.validation import check_X_y, check_array, NotFittedError\n'), (255, 'numpy.mean', 'np.mean', True, 'import numpy as np\n'), (257, 'numpy.unique', 'np.unique', True, 'import numpy as np\n'), (285, 'numpy.array', 'np.array', True, 'import numpy as np\n'), (37, 'sklearn.utils.multiclass.type_of_target', 'type_of_target', False, 'from sklearn.utils.multiclass import check_classification_targets, type_of_target\n'), (42, 'numpy.unique', 'np.unique', True, 'import numpy as np\n'), (43, 'numpy.ones', 'np.ones', True, 'import numpy as np\n'), (121, 'numpy.where', 'np.where', True, 'import numpy as np\n'), (194, 'keras.layers.Dense', 'layers.Dense', False, 'from keras import layers\n'), (48, 'numpy.where', 'np.where', True, 'import numpy as np\n'), (74, 'numpy.where', 'np.where', True, 'import numpy as np\n'), (202, 'tensorflow.keras.optimizers.Adam', 'keras.optimizers.Adam', False, 'from tensorflow import keras\n'), (258, 'numpy.where', 'np.where', True, 'import numpy as np\n'), (260, 'numpy.mean', 'np.mean', True, 'import numpy as np\n'), (50, 'numpy.unique', 'np.unique', True, 'import numpy as np\n'), (52, 'numpy.array', 'np.array', True, 'import numpy as np\n'), (52, 'numpy.sum', 'np.sum', True, 'import numpy as np\n'), (78, 'numpy.array', 'np.array', True, 'import numpy as np\n'), (78, 'numpy.sum', 'np.sum', True, 'import numpy as np\n'), (208, 'tensorflow.keras.callbacks.EarlyStopping', 'keras.callbacks.EarlyStopping', False, 'from tensorflow import keras\n'), (50, 'numpy.where', 'np.where', True, 'import numpy as np\n'), (56, 'numpy.unique', 'np.unique', True, 'import numpy as np\n'), (76, 'numpy.where', 'np.where', True, 'import numpy as np\n')]