#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
@author: Ilker Birbil @ UvA
"""
import warnings
import numpy as np
from typing import List, Optional, Callable, NamedTuple
from sklearn.base import is_classifier, is_regressor
from sklearn.base import RegressorMixin, BaseEstimator, ClassifierMixin
from sklearn.multiclass import OneVsOneClassifier, OneVsRestClassifier, OutputCodeClassifier
from sklearn.neighbors import KDTree
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LinearRegression
from sklearn.preprocessing import StandardScaler
from sklearn.tree import DecisionTreeClassifier, DecisionTreeRegressor
from sklearn.utils.validation import check_X_y, check_array, check_is_fitted
from scipy.stats import mode
from mpi4py import MPI

comm = MPI.COMM_WORLD
number_of_workers=comm.Get_size()
rank = comm.Get_rank()


############################
warnings.formatwarning = lambda msg, *args, **kwargs: f'\nWARNING: \n'+' '.join(str(msg).split())+'\n'
def _LESSwarn(msg, flag=True):
    if (flag):
        warnings.warn(msg)
############################

############################
# Supporting classes

class SklearnEstimator:
    '''
    Dummy base class
    '''
    def fit(self, X: np.array, y: np.array):
        '''
        Dummy fit function
        '''
        raise NotImplementedError('Needs to implement fit(X, y)')

    def predict(self, X0: np.array):
        '''
        Dummy predict function
        '''
        raise NotImplementedError('Needs to implement predict(X, y)')

class LocalModelR(NamedTuple):
    '''
    Auxiliary class to hold the local estimators for regression
    '''
    estimator: SklearnEstimator
    center: np.array

class ReplicationR(NamedTuple):
    '''
    Auxiliary class to hold the replications for regression
    '''
    global_estimator: SklearnEstimator
    local_estimators: List[LocalModelR]

############################

############################
def rbf(data, center, coeff=0.01):
    '''
    RBF kernel - L2 norm
    This is is used as the default distance function in LESS
    '''
    return np.exp(-coeff * np.linalg.norm(np.array(data - center, dtype=float), ord=2, axis=1))

############################

class _LESS(BaseEstimator, SklearnEstimator):
    '''
    The base class for LESSRegressor and LESSClassifier
    '''

    def __init__(self):

        # List to store the replications
        self._replications: Optional[List[ReplicationR]] = None
        # Scaling object used for normalization (StandardScaler)
        self._scobject = None
        # Flag to check whether LESS is fitted 
        self._isfitted = False

    def _set_local_attributes(self):
        '''
        Storing the local variables and checking the given parameters
        '''

        if self.local_estimator is None:
            raise ValueError('LESS does not work without a local estimator.')

        if is_classifier(self.local_estimator):
            _LESSwarn('''
                     LESS might work with local classifiers. 
                     However, we recommend using regressors as the local estimators.
                     ''', self.warnings)

        if (type(self) == LESSRegressor and is_classifier(self.global_estimator)):
            _LESSwarn('''
                     LESSRegressor might work with a global classifier.
                     However, we recommend using a regressor as the global estimator.
                     ''', self.warnings)

        if (type(self) == LESSClassifier and is_regressor(self.global_estimator)):
            _LESSwarn('''
                     LESSClassifier might work with a global regressor. 
                     However, we recommend using a classifier as the global estimator.
                     ''', self.warnings)

        if self.val_size is not None:
            if(self.val_size <= 0.0 or self.val_size >= 1.0):
                raise ValueError('Parameter val_size should be in the interval (0, 1).')

        if self.frac is not None:
            if(self.frac <= 0.0 or self.frac > 1.0):
                raise ValueError('Parameter frac should be in the interval (0, 1].')

        if self.n_replications < 1:
            raise ValueError('The number of replications should greater than equal to one.')

        if self.cluster_method is not None:                       
            if self.frac is not None \
                or self.n_neighbors is not None \
                    or self.n_subsets is not None:
                _LESSwarn('''
                         Parameter cluster_method overrides parameters frac, n_neighbors and n_subsets. \
                         Proceeding with clustering...
                         ''', self.warnings)
                self.frac = None
                self.n_neighbors = None

            # Different numbers of subsets may be generated by the clustering method
            self.n_subsets = []

            if 'n_clusters' in self.cluster_method().get_params().keys():
                if self.cluster_method().get_params()['n_clusters'] == 1:
                    _LESSwarn('''
                             There is only one cluster, so the
                             global estimator is set to none.
                             ''', self.warnings)
                    self.global_estimator = None
                    self.d_normalize = True
                    # If there is also no validation step, then there is
                    # no randomness. So, no need for replications.
                    if (self.val_size is None):
                        _LESSwarn('''
                                 Since validation set is not used, there is no randomness.
                                 Thus, the number of replications is set to one.
                                 ''', self.warnings)
                        self.n_replications = 1
        elif (self.frac is None and
             self.n_neighbors is None and
             self.n_subsets is None):
            self.frac = 0.05
    
    def _check_input(self, len_X: int):
        '''
        Checks whether the input is valid (len_X is the length of input data)
        '''

        if self.cluster_method is None:

            if self.frac is not None:
                self.n_neighbors = int(np.ceil(self.frac * len_X))
                self.n_subsets = int(len_X/self.n_neighbors)

            if self.n_subsets is None:
                self.n_subsets = int(len_X/self.n_neighbors)

            if self.n_neighbors is None:
                self.n_neighbors = int(len_X/self.n_subsets)

            if self.n_neighbors > len_X:
                _LESSwarn('''
                         The number of neighbors is larger than the 
                         number of samples. Setting number of subsets to one.
                         ''', self.warnings)
                self.n_neighbors = len_X
                self.n_subsets = 1

            if self.n_subsets > len_X:
                _LESSwarn('''
                         The number of subsets is larger than the 
                         number of samples. Setting number of neighbors to one.
                         ''', self.warnings)
                self.n_neighbors = 1
                self.n_subsets = len_X

            if self.n_subsets == 1:
                _LESSwarn('''
                         There is only one subset, so the 
                         global estimator is set to none.
                         ''', self.warnings)
                self.global_estimator = None
                self.d_normalize = True
                # If there is also no validation step, then there is
                # no randomness. So, no need for replications.
                if (self.val_size is None):
                    _LESSwarn('''
                             Since validation set is not used, there is no randomness.
                             Thus, the number of replications is set to one.
                             ''', self.warnings)
                    self.n_replications = 1

    def _fit_helper(self,X,y,neighbor_indices_list, Xval = None):
        if rank < ((self.n_subsets) % number_of_workers):
            start = rank * (int((self.n_subsets/number_of_workers))+1)
            stop = start + int((self.n_subsets/number_of_workers))
        else:
            start = (rank * int((self.n_subsets/number_of_workers))) + (self.n_subsets % number_of_workers)
            stop = start + int((self.n_subsets/number_of_workers)) - 1
        my_chunk_len = stop-start+1
        if Xval is None:
            predicts = np.zeros((len(X),my_chunk_len))
            dists = np.zeros((len(X),my_chunk_len))
        else:
            predicts = np.zeros((len(Xval),my_chunk_len))
            dists = np.zeros((len(Xval),my_chunk_len))
            

        local_models: List[LocalModelR] = [None for i in range(my_chunk_len)]
        
        for job_index in range(start,stop+1):
            neighbor_i = job_index
            neighbor_indices = neighbor_indices_list[job_index]
            Xneighbors, yneighbors = X[neighbor_indices], y[neighbor_indices]
            local_center = np.mean(Xneighbors, axis=0)
            if ('random_state' in self.local_estimator().get_params().keys()):
                local_model = self.local_estimator().\
                    set_params(random_state=self.rng_.integers(np.iinfo(np.int16).max)).\
                        fit(Xneighbors, yneighbors)
            else:
                local_model = self.local_estimator().fit(Xneighbors, yneighbors)
            local_models[job_index - start] = (LocalModelR(estimator=local_model, center=local_center))
            if Xval is None:
                predicts[:, job_index - start] = local_model.predict(X)
            else:
                predicts[:, job_index - start] = local_model.predict(Xval)
                
            if(self.distance_function == None):
                if(Xval is None):
                    dists[:, job_index - start] = rbf(X, local_center, \
                        coeff=1.0/np.power(self.n_subsets, 2.0))
                else:
                    dists[:, job_index - start] = rbf(Xval, local_center, \
                        coeff=1.0/np.power(self.n_subsets, 2.0))
            else:
                if(Xval is None):
                    dists[:, job_index - start] = self.distance_function(X, local_center)
                else:
                    dists[:, job_index - start] = self.distance_function(Xval, local_center)
                  

        local_models_gathered = comm.gather(local_models, root=0)
        dists_gathered = comm.gather(dists, root=0)
        predicts_gathered = comm.gather(predicts, root=0)
        if(rank == 0):
            dists_gathered = (np.concatenate(dists_gathered, axis=1))
            local_models_gathered = [localmodel for localmodels in local_models_gathered for localmodel in localmodels]
            predicts_gathered = np.concatenate(predicts_gathered, axis=1)
        return [predicts_gathered, dists_gathered, local_models_gathered]

    def _fitnoval(self, X: np.array, y: np.array):
        '''
        Fit function: All data is used with the global estimator (no validation)
        Tree method is used (no clustering)
        '''

        len_X: int = len(X)
        # Check the validity of the input
        self._check_input(len_X)
        # A nearest neighbor tree is grown for querying
        tree = self.tree_method(X, self.n_subsets)
        self._replications = []
        for _ in range(self.n_replications):
            if rank == 0:
                # Select n_subsets many samples to construct the local sample sets
                sample_indices = self._rng.choice(len_X, size=self.n_subsets)
                # Construct the local sample sets
                _, neighbor_indices_list = np.array(tree.query(X[sample_indices], k=self.n_neighbors), dtype = 'i')
            else:
                neighbor_indices_list = np.zeros([self.n_subsets, self.n_neighbors],dtype='i')
            comm.Bcast(neighbor_indices_list, root=0)
            local_models: List[LocalModelR] = []
            dists = np.zeros((len_X, self.n_subsets))
            predicts = np.zeros((len_X, self.n_subsets))
            [predicts,dists,local_models] = self._fit_helper(X,y,neighbor_indices_list)
            if rank == 0:
                # Normalize the distances from samples to the local subsets
                if self.d_normalize:
                    denom = np.sum(dists, axis=1)
                    denom[denom < 1.0e-8] = 1.0e-8
                    dists = (dists.T/denom).T
            
                if (self.global_estimator != None):
                    if ('random_state' in self.global_estimator().get_params().keys()):
                        global_model = self.global_estimator().\
                            set_params(random_state=self._rng.integers(np.iinfo(np.int16).max)).\
                                fit(dists * predicts, y)
                    else:
                        global_model = self.global_estimator().fit(dists * predicts, y)
                else:
                     global_model = None

                self._replications.append(ReplicationR(global_model, local_models))

        return self

    def _fitval(self, X: np.array, y: np.array):
        '''
        Fit function: (val_size x data) is used for the global estimator (validation)
        Tree method is used (no clustering)
        '''

        self._replications = []
        for i in range(self.n_replications):
            if rank == 0:
                # Split for global estimation
                X_train, X_val, y_train, y_val = train_test_split(X, y,
                    test_size=self.val_size,
                    random_state=self._rng.integers(np.iinfo(np.int16).max))
            else:
                len_x_val = int(len(X) * self.val_size)
                len_y_val = int(len(y) * self.val_size)
                len_x_train = int(len(X) * (1-self.val_size))
                len_y_train = int(len(y) * (1-self.val_size))
                X_train = np.empty((len_x_train,X[0].shape[0]))
                y_train = np.empty(len_y_train)
                X_val = np.empty((len_x_val,X[0].shape[0]))
                y_val = np.empty(len_y_val)

            X_train = comm.bcast(X_train, root=0)
            X_val = comm.bcast(X_val, root=0)
            y_train = comm.bcast(y_train, root=0)
            y_val = comm.bcast(y_val,root=0)  
            len_X_val: int = len(X_val)
            len_X_train: int = len(X_train)
            # Check the validity of the input
            if i == 0:
                self._check_input(len_X_train)
            if rank == 0:
                # A nearest neighbor tree is grown for querying
                tree = self.tree_method(X_train, self.n_subsets)

                # Select n_subsets many samples to construct the local sample sets
                sample_indices = self._rng.choice(len_X_train, size=self.n_subsets)
                # Construct the local sample sets
                _, neighbor_indices_list = np.array(tree.query(X_train[sample_indices], k=self.n_neighbors), dtype='i')
            else:
                neighbor_indices_list = np.zeros([self.n_subsets, self.n_neighbors], dtype = 'i')
            comm.Bcast(neighbor_indices_list, root=0)
            local_models: List[LocalModelR] = []
            dists = np.zeros((len_X_val, self.n_subsets))
            predicts = np.zeros((len_X_val, self.n_subsets))
            [predicts,dists,local_models] = self._fit_helper(X_train,y_train,neighbor_indices_list,X_val)
            if rank == 0:
                # Normalize the distances from samples to the local subsets
                if self.d_normalize:
                    denom = np.sum(dists, axis=1)
                    denom[denom < 1.0e-8] = 1.0e-8
                    dists = (dists.T/denom).T

                if self.global_estimator is not None:
                    if 'random_state' in self.global_estimator().get_params().keys():
                        global_model = self.global_estimator().\
                            set_params(random_state=self._rng.integers(np.iinfo(np.int16).max)).\
                                fit(dists * predicts, y_val)
                    else:
                        global_model = self.global_estimator().fit(dists * predicts, y_val)
                else:
                    global_model = None

                self._replications.append(ReplicationR(global_model, local_models))
        return self

    def _fitnovalc(self, X: np.array, y: np.array):
        '''
        Fit function: All data is used for the global estimator (no validation)
        Clustering is used (no tree method)
        '''

        len_X: int = len(X)
        # Check the validity of the input
        self._check_input(len_X)
        if ('random_state' not in self.cluster_method_().get_params().keys()): 
                warnings.warn('Clustering method is not random, so there is \
                    no need for replications, unless validaton set is used. \
                    Note that lack of replications may increase the variance.') 
                if(rank == 0):
                     cluster_fit = self.cluster_method().fit(X)
                     self.n_replications_ = 1
                     cluster_fit_labels = cluster_fit.labels_
                     cluster_fit_centers = cluster_fit.cluster_centers_
                     n_clusters = cluster_fit_labels.shape[0]
                     n_features = cluster_fit_labels.shape[1]
                else:
                     n_clusters = comm.bcast(n_clusters, root=0)
                     n_features = comm.bcast(n_features, root=0)
                     self.n_replications_ = 1
                     cluster_fit_labels = np.zeros([n_clusters, n_features], dtype = 'i')
                     cluster_fit_centers = np.zeros([n_clusters, n_features], dtype = 'd')
                cluster_fit_labels = comm.bcast(cluster_fit_labels, root=0) 
                cluster_fit_centers = comm.bcast(cluster_fit_centers, root=0) 
        self.replications_ = []
        for i in range(self.n_replications_):
            if (self.n_replications_ > 1):
                if (rank == 0):
                    cluster_fit = self.cluster_method(n_jobs=1).\
                        set_params(random_state=self.rng_.integers(np.iinfo(np.int16).max)).\
                             fit(X)
                    cluster_fit_labels = cluster_fit.labels_
                    cluster_fit_centers = cluster_fit.cluster_centers_
                    n_clusters = cluster_fit_centers.shape[0]
                    n_features = cluster_fit_centers.shape[1]
                else:
                     n_clusters = 0
                     n_features = 0
                     use_cluster_centers = False
                     cluster_fit_labels = np.zeros([n_clusters, n_features], dtype = 'i')
                     cluster_fit_centers = np.zeros([n_clusters, n_features], dtype = 'd')
                    
                n_clusters = comm.bcast(n_clusters, root=0)
                n_features = comm.bcast(n_features, root=0)
                cluster_fit_labels = comm.bcast(cluster_fit_labels, root=0) 
                cluster_fit_centers = comm.bcast(cluster_fit_centers, root=0) 
            # Some clustering methods may find less number of
            # clusters than requested 'n_clusters'
            self.n_subsets_.append(len(np.unique(cluster_fit_labels)))
            n_subsets = self.n_subsets_[i]
            
            local_models: List[LocalModelR] = []
            dists = np.zeros((len_X, n_subsets))            
            predicts = np.zeros((len_X, n_subsets))
            if rank == 0:            
                if (hasattr(cluster_fit, 'cluster_centers_')):
                     use_cluster_centers = True
                else:
                     use_cluster_centers = False
            use_cluster_centers = comm.bcast(use_cluster_centers, root=0)
            [predicts,dists,local_models] = self._fit_helperc(X,y,cluster_fit_labels,cluster_fit_centers,use_cluster_centers,i)

            if rank == 0:
                # Normalize the distances from samples to the local subsets
                if self.d_normalize:
                    denom = np.sum(dists, axis=1)
                    denom[denom < 1.0e-8] = 1.0e-8
                    dists = (dists.T/denom).T

                if self.global_estimator is not None:
                    if 'random_state' in self.global_estimator().get_params().keys():
                        global_model = self.global_estimator().\
                            set_params(random_state=self._rng.integers(np.iinfo(np.int16).max)).\
                            fit(dists * predicts, y)
                    else:
                        global_model = self.global_estimator().fit(dists * predicts, y)
                else:
                    global_model = None

                self._replications.append(ReplicationR(global_model, local_models))

        return self

    def _fitvalc(self, X: np.array, y: np.array):
        '''
        Fit function: (val_size x data) is used for the global estimator (validation)
        Clustering is used (no tree method)
        '''

        self.replications_ = []
        for i in range(self.n_replications_):
            if rank ==0:
                # Split for global estimation
                X_train, X_val, y_train, y_val = train_test_split(X, y,
                    test_size=self.val_size_,
                    random_state=self.rng_.integers(np.iinfo(np.int16).max))
            else:
                len_x_val = int(len(X) * self.val_size)
                len_y_val = int(len(y) * self.val_size)
                len_x_train = int(len(X) * (1-self.val_size))
                len_y_train = int(len(y) * (1-self.val_size))
                X_train = np.empty((len_x_train,X[0].shape[0]))
                y_train = np.empty(len_y_train)
                X_val = np.empty((len_x_val,X[0].shape[0]))
                y_val = np.empty(len_y_val)
            
            X_train = comm.bcast(X_train, root=0)
            X_val = comm.bcast(X_val, root=0)
            y_train = comm.bcast(y_train, root=0)
            y_val = comm.bcast(y_val,root=0)  
            len_X_val: int = len(X_val)
            len_X_train: int = len(X_train)
            # Check the validity of the input
            if (i == 0):
                use_cluster_centers = False
                self._check_input(len_X_train)

 
            if(rank == 0):
                 if 'random_state' not in self.cluster_method().get_params().keys():
                     cluster_fit = self.cluster_method().fit(X_train)
                 else:
                     cluster_fit = self.cluster_method().\
                         set_params(random_state=self._rng.integers(np.iinfo(np.int16).max)).\
                             fit(X_train)
                 cluster_fit_labels = cluster_fit.labels
                 cluster_fit_centers = cluster_fit.cluster_centers
                 n_clusters = cluster_fit_centers.shape[0]
                 n_features = cluster_fit_centers.shape[1]
            else:
                 n_clusters = 0
                 n_features = 0
                 cluster_fit_labels = np.zeros([n_clusters, n_features], dtype = 'i')
                 cluster_fit_centers = np.zeros([n_clusters, n_features], dtype = 'd')
            n_clusters = comm.bcast(n_clusters, root=0)
            n_features = comm.bcast(n_features, root=0)
            cluster_fit_labels = comm.bcast(cluster_fit_labels, root=0) 
            cluster_fit_centers = comm.bcast(cluster_fit_centers, root=0) 
            if (rank == 0):
                if (i==1):
                    if (hasattr(cluster_fit, 'cluster_centers')):
                        use_cluster_centers = True
                    else:
                        use_cluster_centers = False
            use_cluster_centers = comm.bcast(use_cluster_centers, root=0)  
            # Since each replication returns
            self.n_subsets_.append(len(np.unique(cluster_fit_labels)))
            n_subsets = self.n_subsets[i]
            
            local_models: List[LocalModelR] = []
            dists = np.zeros((len_X_val, n_subsets))            
            predicts = np.zeros((len_X_val, n_subsets))
                
            [predicts,dists,local_models] = self._fit_helperc(X_train,y_train,cluster_fit_labels,cluster_fit_centers,use_cluster_centers,i,X_val)
            if (rank == 0):
                # Normalize the distances from samples to the local subsets
                if self.d_normalize:
                    denom = np.sum(dists, axis=1)
                    denom[denom < 1.0e-8] = 1.0e-8
                    dists = (dists.T/denom).T 
                if (self.global_estimator_ != None):
                    if ('random_state' in self.global_estimator().get_params().keys()):
                        global_model = self.global_estimator().\
                            set_params(random_state=self._rng.integers(np.iinfo(np.int16).max)).\
                                fit(dists * predicts, y_val)
                    else:
                        global_model = self.global_estimator().fit(dists * predicts, y_val)
                else:
                    global_model = None

                self.replications_.append(ReplicationR(global_model, local_models))

        return self

    def get_n_subsets(self):
        '''
        Auxiliary function returning the number of subsets
        '''

        if not self._isfitted:
            _LESSwarn('''
                     You need to fit LESS first.
                     ''', self.warnings)

        return self.n_subsets

    def get_n_neighbors(self):
        '''
        Auxiliary function returning the number of neighbors
        '''

        if self.cluster_method is not None:
            _LESSwarn('''
                     Number of neighbors is not fixed when clustering is used.
                     ''', self.warnings)
        elif not self._isfitted:
            _LESSwarn('''
                     You need to fit LESS first.
                     ''', self.warnings)

        return self.n_neighbors

    def get_frac(self):
        '''
        Auxiliary function returning the percentage of samples used to set the number of neighbors
        '''

        # Fraction is set to None only if clustering method is given
        if self.cluster_method is not None:
            _LESSwarn('''
                     Parameter frac is not set when clustering is used.
                     ''', self.warnings)

        return self.frac

    def get_n_replications(self):
        '''
        Auxiliary function returning the number of replications
        '''

        return self.n_replications

    def get_d_normalize(self):
        '''
        Auxiliary function returning the flag for normalization
        '''

        return self.d_normalize

    def get_scaling(self):
        '''
        Auxiliary function returning the flag for scaling
        '''

        return self.scaling
    
    def get_val_size(self):
        '''
        Auxiliary function returning the validation set size
        '''

        return self.val_size

    def get_random_state(self):
        '''
        Auxiliary function returning the random seed
        '''

        return self.random_state

class LESSClassifier(_LESS, ClassifierMixin):
    '''
    Classifier for Learning with Subset Selection (LESS)

    This is a wrapper that calls the multiclass strategies, like one-vs-rest, 
    by using an auxiliary binary classifer for LESS (_LBC)

    Parameters
    ----------
        frac: fraction of total samples used for the number of neighbors (default is 0.05)
        n_neighbors : number of neighbors (default is None)
        n_subsets : number of subsets (default is None)
        n_replications : number of replications (default is 50)
        d_normalize : distance normalization (default is True)
        val_size: percentage of samples used for validation (default is None - no validation)
        random_state: initialization of the random seed (default is None)
        tree_method : method used for constructing the nearest neighbor tree,
                e.g., sklearn.neighbors.KDTree (default) or sklearn.neighbors.BallTree
        cluster_method : method used for clustering the subsets,
                e.g., sklearn.cluster.KMeans, sklearn.cluster.SpectralClustering (default is None)
        local_estimator : estimator for training the local models
                (default is LinearRegression)
        global_estimator : estimator for training the global model
                (default is DecisionTreeClassifier)
        distance_function : distance function evaluating the distance from a subset to a sample,
                e.g., df(subset, sample) which returns a vector of distances
                (default is RBF(subset, sample, 1.0/n_subsets^2))
        scaling: flag to normalize the input data (default is True)
        warnings : flag to turn on (True) or off (False) the warnings (default is True)
        multiclass : available strategies are 'ovr' (one-vs-rest, default), 
                'ovo' (one-vs-one), 'occ' (output-code-classifier)

    Recommendation
    --------------
    Default implementation of LESS uses Euclidean distances with radial basis function.
    Therefore, it is a good idea to scale the input data before fitting. This can be done by
    setting the parameter 'scaling' to True (the default value) or preprocessing the data
    as follows:

    >>> from sklearn.preprocessing import StandardScaler
    >>> SC = StandardarScaler()
    >>> X_train = SC.fit_transform(X_train)
    >>> X_test = SC.transform(X_test)

    '''

    def __init__(self, frac=None, n_neighbors=None, n_subsets=None,
                n_replications=20, d_normalize=True, val_size=None, random_state=None,
                tree_method=lambda data, n_subsets: KDTree(data, n_subsets),
                cluster_method=None,
                local_estimator=lambda: LinearRegression(),
                global_estimator=lambda: DecisionTreeClassifier(),
                distance_function: Callable[[np.array, np.array], np.array]=None,
                scaling=True, warnings=True, multiclass='ovr'):

        self.local_estimator = local_estimator
        self.global_estimator = global_estimator
        self.tree_method = tree_method
        self.cluster_method = cluster_method
        self.distance_function = distance_function
        self.frac = frac
        self.n_neighbors = n_neighbors
        self.n_subsets = n_subsets
        self.n_replications = n_replications
        self.d_normalize = d_normalize
        self.val_size = val_size
        self.random_state = random_state
        self._bclassifier = None
        self._strategy = None
        self.scaling = scaling
        self.warnings = warnings
        self.multiclass = multiclass

        class _LESSBC(_LESS):
            '''
            Auxiliary binary classifier for Learning with Subset Selection (LESS)
            
            NOTE: There is no scaling option in this class

            '''
            def __init__(self, frac=None, n_neighbors=None, n_subsets=None,
                        n_replications=20, d_normalize=True, val_size=None, random_state=None,
                        tree_method=lambda data, n_subsets: KDTree(data, n_subsets),
                        cluster_method=None,
                        local_estimator=lambda: LinearRegression(),
                        global_estimator=lambda: DecisionTreeClassifier(),
                        distance_function: Callable[[np.array, np.array], np.array]=None,
                        warnings=True):

                self.local_estimator = local_estimator
                self.global_estimator = global_estimator
                self.tree_method = tree_method
                self.cluster_method = cluster_method
                self.distance_function = distance_function
                self.frac = frac
                self.n_neighbors = n_neighbors
                self.n_subsets = n_subsets
                self.n_replications = n_replications
                self.d_normalize = d_normalize
                self.val_size = val_size
                self.random_state = random_state
                self._rng = np.random.default_rng(self.random_state)
                self.warnings = warnings

            def fit(self, X: np.array, y: np.array):
                '''
                Dummy fit function that calls the proper method according to 
                validation and clustering parameters
                
                Options are:
                - Default fitting (no validation set, no clustering)
                - Fitting with validation set (no clustering)
                - Fitting with clustering (no) validation set)
                - Fitting with validation set and clustering
                '''

                # Check that X and y have correct shape
                X, y = check_X_y(X, y)       

                # Original labels
                self._yorg = np.unique(y)

                if len(self._yorg) != 2:
                    raise ValueError('LESSBinaryClassifier works only with two labels. \
                                    Please try LESSClassifier.')

                # Convert to binary labels
                ymin1 = y == self._yorg[0]
                ypls1 = y == self._yorg[1]
                y[ymin1] = -1
                y[ypls1] = 1

                self._set_local_attributes()

                if self.val_size is not None:
                    # Validation set is used for
                    # global estimation
                    if self.cluster_method is None:
                        self._fitval(X, y)
                    else:
                        self._fitvalc(X, y)
                else:
                    # Validation set is not used for
                    # global estimation
                    if self.cluster_method is None:
                        self._fitnoval(X, y)
                    else:
                        self._fitnovalc(X, y)

                # Convert to original labels
                ymin1 = y == -1
                ypls1 = y == 1
                y[ymin1] = self._yorg[0]        
                y[ypls1] = self._yorg[1]

                self._isfitted = True

                return self

            def predict(self, X0: np.array):
                '''
                Predictions are evaluated for the test samples in X0
                '''

                check_is_fitted(self, attributes='_isfitted')
                # Input validation
                X0 = check_array(X0)

                len_X0: int = len(X0)
                yhat = np.zeros((len_X0, self.n_replications))
                for i in range(self.n_replications):
                    # Get the fitted global and local estimators
                    global_model = self._replications[i].global_estimator
                    local_models = self._replications[i].local_estimators
                    if self.cluster_method is None:
                        n_subsets = self.n_subsets
                    else:
                        n_subsets = self.n_subsets[i]
                    predicts = np.zeros((len_X0, n_subsets))
                    dists = np.zeros((len_X0, n_subsets))
                    for j in range(n_subsets):
                        local_center = local_models[j].center
                        local_model = local_models[j].estimator
                        predicts[:, j] = local_model.predict(X0)

                        if self.distance_function is None:
                            dists[:, j] = rbf(X0, local_center, \
                                coeff=1.0/np.power(n_subsets, 2.0))
                        else:
                            dists[:, j] = self.distance_function(X0, local_center)

                    # Normalize the distances from samples to the local subsets
                    if self.d_normalize:
                        denom = np.sum(dists, axis=1)
                        denom[denom < 1.0e-8] = 1.0e-8
                        dists = (dists.T/denom).T

                    if global_model is not None:
                        yhat[:, i] = global_model.predict(dists * predicts)
                    else:
                        rowsums = np.sum(dists * predicts, axis=1)
                        yhat[rowsums < 0, i] = -1
                        yhat[rowsums >= 0, i] = 1

                yhat = mode(yhat.astype(int), axis=1).mode.reshape(1, -1)[0]

                # Convert to original labels
                ymin1 = yhat == -1
                ypls1 = yhat == 1        
                yhat[ymin1] = self._yorg[0]
                yhat[ypls1] = self._yorg[1]

                return yhat

            def predict_proba(self, X0: np.array):
                '''
                Prediction probabilities are evaluated for the test samples in X0
                '''

                check_is_fitted(self, attributes='_isfitted')
                # Input validation
                X0 = check_array(X0)

                len_X0: int = len(X0)
                yhat = np.zeros((len_X0, self.n_replications), dtype=np.int)
                predprobs = np.zeros((len_X0, 2), dtype=np.float16)
                for i in range(self.n_replications):
                    # Get the fitted global and local estimators
                    global_model = self._replications[i].global_estimator
                    local_models = self._replications[i].local_estimators
                    if self.cluster_method is None:
                        n_subsets = self.n_subsets
                    else:
                        n_subsets = self.n_subsets[i]
                    predicts = np.zeros((len_X0, n_subsets))
                    dists = np.zeros((len_X0, n_subsets))
                    for j in range(n_subsets):
                        local_center = local_models[j].center
                        local_model = local_models[j].estimator
                        predicts[:, j] = local_model.predict(X0)

                        if self.distance_function is None:
                            dists[:, j] = rbf(X0, local_center, \
                                coeff=1.0/np.power(n_subsets, 2.0))
                        else:
                            dists[:, j] = self.distance_function(X0, local_center)

                    # Normalize the distances from samples to the local subsets
                    if self.d_normalize:
                        denom = np.sum(dists, axis=1)
                        denom[denom < 1.0e-8] = 1.0e-8
                        dists = (dists.T/denom).T

                    if global_model is not None:
                        yhat[:, i] = global_model.predict(dists * predicts)
                        # Convert to 0-1
                        yhat[:, i] = (yhat[:, i] + 1)/2
                    else:
                        rowsums = np.sum(dists * predicts, axis=1)
                        yhat[rowsums < 0, i] = 0
                        yhat[rowsums >= 0, i] = 1

                md, cnt = mode(yhat, axis=1)
                yhat = md.reshape(1, -1)[0]
                cnt = cnt.reshape(1, -1)[0]
                yhat0 = yhat==0
                yhat1 = yhat==1
                predprobs[yhat0, 0] = cnt[yhat0]
                predprobs[yhat0, 1] = self.n_replications - cnt[yhat0]
                predprobs[yhat1, 1] = cnt[yhat1]
                predprobs[yhat1, 0] = self.n_replications - cnt[yhat1]

                predprobs /= self.n_replications

                return predprobs

        self._bclassifier = _LESSBC(frac=self.frac, n_neighbors=self.n_neighbors,
                                    n_subsets=self.n_subsets,
                                    n_replications=self.n_replications,
                                    d_normalize=self.d_normalize,
                                    val_size=self.val_size,
                                    random_state=self.random_state,
                                    tree_method=self.tree_method,
                                    cluster_method=self.cluster_method,
                                    local_estimator=self.local_estimator,
                                    global_estimator=self.global_estimator,
                                    distance_function=self.distance_function,
                                    warnings=self.warnings)

    def fit(self, X: np.array, y: np.array):
        '''
        Dummy fit function that calls the fit method of the multiclass strategy 'one-vs-rest'
        '''
        if (self.scaling):
            self._scobject = StandardScaler()
            X = self._scobject.fit_transform(X)

        n_classes = len(np.unique(y))
        
        self._set_strategy(n_classes)

        self._strategy.fit(X, y)
       
        self._update_params(self._strategy.estimators_[0], n_classes)
        
        self._isfitted = True

        return self

    def predict(self, X0: np.array):
        '''
        Dummy predict function that calls the predict method of the multiclass strategy 'one-vs-rest'
        '''

        if (self.scaling):
            X0 = self._scobject.transform(X0)

        return self._strategy.predict(X0)
    
    def _set_strategy(self, n_classes):
        '''
        Auxiliary function to set the selected the strategy
        '''
        
        if (n_classes == 2):
            self._strategy = OneVsRestClassifier(self._bclassifier)
        elif (self.multiclass == 'ovr'):
            self._strategy = OneVsRestClassifier(self._bclassifier)
        elif (self.multiclass == 'ovo'):
            self._strategy = OneVsOneClassifier(self._bclassifier)
        elif (self.multiclass == 'occ'):
            self._strategy = OutputCodeClassifier(self._bclassifier)
        else:
            self._strategy = OneVsRestClassifier(self._bclassifier)
            _LESSwarn('''
                      LESSClassifier works only with one of the following options:
                      (1) 'ovr' : OneVsRestClassifier (default),
                      (2) 'ovo' : OneVsOneClassifier,
                      (3) 'occ' : OutputCodeClassifier,
                      (see sklearn.multiclass for details).

                      Switching to 'ovr' ...
                      ''', self.warnings)

    def _update_params(self, firstestimator, n_classes):
        '''
        Parameters of the wrapper class are updated, since the functions
        _set_local_attributes and _check_input may alter the following parameters
        '''
        
        self.global_estimator = firstestimator.global_estimator
        self.frac = firstestimator.get_frac()
        self.n_neighbors = firstestimator.get_n_neighbors()
        self.n_subsets = firstestimator.get_n_subsets()
        self.n_replications = firstestimator.get_n_replications()
        self.d_normalize = firstestimator.get_d_normalize()
        # Replications are stored only if it is a binary classification problem
        # Otherwise, there are multiple binary classifiers, and hence, multiple replications
        if n_classes == 2:
            self._replications = firstestimator._replications        

class LESSRegressor(_LESS, RegressorMixin):
    '''
    Regressor for Learning with Subset Selection (LESS)

    Parameters
    ----------
        frac: fraction of total samples used for the number of neighbors (default is 0.05)
        n_neighbors : number of neighbors (default is None)
        n_subsets : number of subsets (default is None)
        n_replications : number of replications (default is 50)
        d_normalize : distance normalization (default is True)
        val_size: percentage of samples used for validation (default is None - no validation)
        random_state: initialization of the random seed (default is None)
        tree_method : method used for constructing the nearest neighbor tree,
                e.g., sklearn.neighbors.KDTree (default) or sklearn.neighbors.BallTree
        cluster_method : method used for clustering the subsets,
                e.g., sklearn.cluster.KMeans, sklearn.cluster.SpectralClustering (default is None)
        local_estimator : estimator for training the local models (default is LinearRegression)
        global_estimator : estimator for training the global model (default is DecisionTreeRegressor)
        distance_function : distance function evaluating the distance from a subset to a sample,
                e.g., df(subset, sample) which returns a vector of distances
                (default is RBF(subset, sample, 1.0/n_subsets^2))
        scaling: flag to normalize the input data (default is True)
        warnings : flag to turn on (True) or off (False) the warnings (default is True)

    Recommendation
    --------------
    Default implementation of LESS uses Euclidean distances with radial basis function.
    Therefore, it is a good idea to scale the input data before fitting. This can be done by
    setting the parameter 'scaling' to True (the default value) or preprocessing the data
    as follows:

    >>> from sklearn.preprocessing import StandardScaler
    >>> SC = StandardarScaler()
    >>> X_train = SC.fit_transform(X_train)
    >>> X_test = SC.transform(X_test)

    '''

    def __init__(self, frac=None, n_neighbors=None, n_subsets=None,
                 n_replications=20, d_normalize=True, val_size=None, random_state=None,
                 tree_method=lambda data, n_subsets: KDTree(data, n_subsets),
                 cluster_method=None,
                 local_estimator=lambda: LinearRegression(),
                 global_estimator=lambda: DecisionTreeRegressor(),
                 distance_function: Callable[[np.array, np.array], np.array]=None,
                 scaling=True, warnings=True):

        self.local_estimator = local_estimator
        self.global_estimator = global_estimator
        self.tree_method = tree_method
        self.cluster_method = cluster_method
        self.distance_function = distance_function
        self.frac = frac
        self.n_neighbors = n_neighbors
        self.n_subsets = n_subsets
        self.n_replications = n_replications
        self.d_normalize = d_normalize
        self.val_size = val_size
        self.random_state = random_state
        self._rng = np.random.default_rng(self.random_state)
        self.scaling = scaling
        self.warnings = warnings

    def fit(self, X: np.array, y: np.array):
        '''
        Dummy fit function that calls the proper method according to 
        validation and clustering parameters
        
        Options are:
          - Default fitting (no validation set, no clustering)
          - Fitting with validation set (no clustering)
          - Fitting with clustering (no) validation set)
          - Fitting with validation set and clustering
        '''

        # Check that X and y have correct shape
        X, y = check_X_y(X, y)

        self._set_local_attributes()
        
        if (self.scaling):
            self._scobject = StandardScaler()
            X = self._scobject.fit_transform(X)

        if self.val_size is not None:
            # Validation set is not used for
            # global estimation
            if self.cluster_method is None:
                self._fitval(X, y)
            else:
                self._fitvalc(X, y)
        else:
            # Validation set is used for
            # global estimation
            if self.cluster_method is None:
                self._fitnoval(X, y)
            else:
                self._fitnovalc(X, y)

        self._isfitted = True

        return self

    def predict(self, X0: np.array):
        '''
        Predictions are evaluated for the test samples in X0
        '''

        check_is_fitted(self, attributes='_isfitted')
        # Input validation
        X0 = check_array(X0)

        if (self.scaling):
            X0 = self._scobject.transform(X0)

        len_X0: int = len(X0)
        yhat = np.zeros(len_X0)
        for i in range(self.n_replications):
            # Get the fitted global and local estimators
            global_model = self._replications[i].global_estimator
            local_models = self._replications[i].local_estimators
            if self.cluster_method is None:
                n_subsets = self.n_subsets
            else:
                n_subsets = self.n_subsets[i]
            predicts = np.zeros((len_X0, n_subsets))
            dists = np.zeros((len_X0, n_subsets))
            for j in range(n_subsets):
                local_center = local_models[j].center
                local_model = local_models[j].estimator
                predicts[:, j] = local_model.predict(X0)

                if self.distance_function is None:
                    dists[:, j] = rbf(X0, local_center, \
                        coeff=1.0/np.power(n_subsets, 2.0))
                else:
                    dists[:, j] = self.distance_function(X0, local_center)

            # Normalize the distances from samples to the local subsets
            if self.d_normalize:
                denom = np.sum(dists, axis=1)
                denom[denom < 1.0e-8] = 1.0e-8
                dists = (dists.T/denom).T

            if global_model is not None:
                yhat += global_model.predict(dists * predicts)
            else:
                yhat += np.sum(dists * predicts, axis=1)

        yhat = yhat/self.n_replications

        return yhat
