'''Wraps algorithms with SklearnAlgorithm and registers them.'''

# Copyright (c) 2023 Carnegie Mellon University
# This code is subject to the license terms contained in the LICENSE file.

from typing import Any, Dict, List

from .sklearn_algorithm import SklearnAlgorithm
from ...algorithms.impl.algorithm import MemoryAlgorithmCatalog
from ...problem_def.task import DataType, TaskType
from ...searcher.param_range import ParamRange
from ...wrangler.constants import RangeMethod


ALGORITHMS: List[Dict[str, Any]] = [
    # Classification Algorithms
    {
        'name': 'sklearn.linear_model.LogisticRegression',
        'tasks': [TaskType.BINARY_CLASSIFICATION,
                  TaskType.MULTICLASS_CLASSIFICATION],
        'data_types': [DataType.TABULAR],
        'tags': {
            'for_tests': ['true'],
            'supports_random_seed': ['true']
        },
        # Hyperparams
        'penalty': 'l2',
        'dual': False,
        'tol': 1e-4,
        'C': ParamRange(method=RangeMethod.LIST,
                        prange=[0.1, 1.0, 10.0, 100.0],
                        default=1.0),
        'fit_intercept': True,
        'intercept_scaling': 1.0,
        'class_weight': ParamRange(method=RangeMethod.LIST,
                                   prange=['balanced', None]),
        'random_state': None,
        'solver': 'lbfgs',
        'max_iter': 1000,
        'multi_class': 'auto',
        'verbose': 0,
        'warm_start': False,
        'n_jobs': None,
        'l1_ratio': None
    },
    {
        'name': 'sklearn.ensemble.RandomForestClassifier',
        'tasks': [TaskType.BINARY_CLASSIFICATION,
                  TaskType.MULTICLASS_CLASSIFICATION],
        'data_types': [DataType.TABULAR],
        'tags': {'supports_random_seed': ['true']},
        # Hyperparams
        'n_estimators': 100,
        'criterion': 'gini',
        'max_depth': ParamRange(method=RangeMethod.LIST, prange=[8, 10, 15, None]),
        'min_samples_split': ParamRange(method=RangeMethod.LIST, prange=[2, 5, 10], default=2),
        'min_samples_leaf': 1,
        'min_weight_fraction_leaf': 0.0,
        'max_features': 'sqrt',
        'max_leaf_nodes': None,
        'min_impurity_decrease': 0.0,
        'bootstrap': True,
        'oob_score': False,
        'n_jobs': None,
        'random_state': None,
        'verbose': 0,
        'warm_start': False,
        'class_weight': ParamRange(method=RangeMethod.LIST, prange=['balanced', None]),
        'ccp_alpha': 0.0,
        'max_samples': None,
    },

    {
        'name': 'sklearn.svm.SVC',
        'tasks': [TaskType.BINARY_CLASSIFICATION,
                  TaskType.MULTICLASS_CLASSIFICATION],
        'data_types': [DataType.TABULAR],
        'tags': {
            'has_predict_proba': ['false'],
            'for_tests': ['true'],
            'supports_random_seed': ['true']},
        # Hyperparams
        'C': ParamRange(method=RangeMethod.LIST,
                        prange=[0.01, 0.1, 1.0, 10.0, 100.0],
                        default=1.0),
        'kernel': 'rbf',
        'degree': 3,
        'gamma': 'scale',
        'coef0': 0.0,
        'shrinking': True,
        'probability': False,
        'tol': 1e-3,
        'cache_size': 200,
        'class_weight': ParamRange(method=RangeMethod.LIST,
                                   prange=['balanced', None]),
        'verbose': False,
        'max_iter': -1,
        'decision_function_shape': 'ovr',
        'break_ties': False,
        'random_state': None
    },
    {
        'name': 'sklearn.naive_bayes.MultinomialNB',
        'tasks': [TaskType.BINARY_CLASSIFICATION,
                  TaskType.MULTICLASS_CLASSIFICATION],
        'data_types': [DataType.TABULAR],
        # Hyperparams
        'alpha': 1.0,
        'force_alpha': True,
        'fit_prior': True,
        'class_prior': None
    },
    {
        'name': 'sklearn.discriminant_analysis.LinearDiscriminantAnalysis',
        'tasks': [TaskType.BINARY_CLASSIFICATION,
                  TaskType.MULTICLASS_CLASSIFICATION],
        'data_types': [DataType.TABULAR],
        # Hyperparams
        'solver': 'svd',
        'shrinkage': None,
        'priors': None,
        'n_components': None,
        'store_covariance': False,
        'tol': 0.0001,
        'covariance_estimator': None
    },
    {
        'name': 'sklearn.discriminant_analysis.QuadraticDiscriminantAnalysis',
        'tasks': [TaskType.BINARY_CLASSIFICATION,
                  TaskType.MULTICLASS_CLASSIFICATION],
        'data_types': [DataType.TABULAR],
        # Hyperparams
        'priors': None,
        'reg_param': 0.0,
        'store_covariance': False,
        'tol': 0.0001
    },
    {
        'name': 'sklearn.ensemble.AdaBoostClassifier',
        'tasks': [TaskType.BINARY_CLASSIFICATION,
                  TaskType.MULTICLASS_CLASSIFICATION],
        'data_types': [DataType.TABULAR],
        'tags': {'supports_random_seed': ['true']},
        # Hyperparams
        'estimator': None,
        'n_estimators': 50,
        'learning_rate': 1.0,
        'random_state': None,
    },
    {
        'name': 'sklearn.svm.LinearSVC',
        'tasks': [TaskType.BINARY_CLASSIFICATION,
                  TaskType.MULTICLASS_CLASSIFICATION],
        'data_types': [DataType.TABULAR],
        'tags': {
            'has_predict_proba': ['false'],
            'supports_random_seed': ['true']},
        # Hyperparams
        'penalty': 'l2',
        'loss': 'squared_hinge',
        'dual': True,
        'tol': 0.0001,
        'C': ParamRange(method=RangeMethod.LIST,
                        prange=[0.01, 0.1, 1.0, 10.0, 100.0],
                        default=1.0),
        'multi_class': 'ovr',
        'fit_intercept': True,
        'intercept_scaling': 1,
        'class_weight': ParamRange(method=RangeMethod.LIST,
                                   prange=['balanced', None]),
        'verbose': 0,
        'random_state': None,
        'max_iter': 1000,
    },
    {
        'name': 'sklearn.tree.ExtraTreeClassifier',
        'tasks': [TaskType.BINARY_CLASSIFICATION,
                  TaskType.MULTICLASS_CLASSIFICATION],
        'data_types': [DataType.TABULAR],
        'tags': {'supports_random_seed': ['true']},
        # Hyperparams
        'criterion': 'gini',
        'splitter': 'random',
        'max_depth': ParamRange(method=RangeMethod.LIST,
                                prange=[8, 10, 15, None]),
        'min_samples_split': ParamRange(method=RangeMethod.LIST,
                                        prange=[2, 5, 10],
                                        default=2),
        'min_samples_leaf': 1,
        'min_weight_fraction_leaf': 0.0,
        'max_features': 'sqrt',
        'random_state': None,
        'max_leaf_nodes': None,
        'min_impurity_decrease': 0.0,
        'class_weight': ParamRange(method=RangeMethod.LIST, prange=['balanced', None]),
        'ccp_alpha': 0.0,
    },
    {
        'name': 'sklearn.ensemble.BaggingClassifier',
        'tasks': [TaskType.BINARY_CLASSIFICATION,
                  TaskType.MULTICLASS_CLASSIFICATION],
        'data_types': [DataType.TABULAR],
        'tags': {'supports_random_seed': ['true']},
        # Hyperparams
        'estimator': None,
        'n_estimators': 10,
        'max_samples': 1.0,
        'max_features': 1.0,
        'bootstrap': True,
        'bootstrap_features': False,
        'oob_score': False,
        'warm_start': False,
        'n_jobs': None,
        'random_state': None,
        'verbose': 0,
    },
    {
        'name': 'sklearn.linear_model.PassiveAggressiveClassifier',
        'tasks': [TaskType.BINARY_CLASSIFICATION,
                  TaskType.MULTICLASS_CLASSIFICATION],
        'data_types': [DataType.TABULAR],
        'tags': {
            'has_predict_proba': ['false'],
            'supports_random_seed': ['true']},
        # Hyperparams
        'C': 1.0,
        'fit_intercept': True,
        'max_iter': 1000,
        'tol': 0.001,
        'early_stopping': False,
        'validation_fraction': 0.1,
        'n_iter_no_change': 5,
        'shuffle': True,
        'verbose': 0,
        'loss': 'hinge',
        'n_jobs': None,
        'random_state': None,
        'warm_start': False,
        'class_weight': None,
        'average': False
    },
    {
        'name': 'sklearn.neural_network.MLPClassifier',
        'tasks': [TaskType.BINARY_CLASSIFICATION,
                  TaskType.MULTICLASS_CLASSIFICATION],
        'data_types': [DataType.TABULAR],
        'tags': {'supports_random_seed': ['true']},
        # Hyperparams
        'hidden_layer_sizes': (100,),
        'activation': 'relu',
        'solver': 'adam',
        'alpha': 0.0001,
        'batch_size': 'auto',
        'learning_rate': 'constant',
        'learning_rate_init': 0.001,
        'power_t': 0.5,
        'max_iter': 200,
        'shuffle': True,
        'random_state': None,
        'tol': 0.0001,
        'verbose': False,
        'warm_start': False,
        'momentum': 0.9,
        'nesterovs_momentum': True,
        'early_stopping': False,
        'validation_fraction': 0.1,
        'beta_1': 0.9,
        'beta_2': 0.999,
        'epsilon': 1e-08,
        'n_iter_no_change': 10,
        'max_fun': 15000,
    },
    {
        'name': 'sklearn.ensemble.GradientBoostingClassifier',
        'tasks': [TaskType.BINARY_CLASSIFICATION,
                  TaskType.MULTICLASS_CLASSIFICATION],
        'data_types': [DataType.TABULAR],
        'tags': {'supports_random_seed': ['true']},
        # Hyperparams
        'loss': 'log_loss',
        'learning_rate': 0.1,
        'n_estimators': 100,
        'subsample': 1.0,
        'criterion': 'friedman_mse',
        'min_samples_split': ParamRange(method=RangeMethod.LIST,
                                        prange=[2, 5, 10],
                                        default=2),
        'min_samples_leaf': ParamRange(method=RangeMethod.LIST,
                                       prange=[1, 2, 5],
                                       default=1),
        'min_weight_fraction_leaf': 0.0,
        'max_depth': ParamRange(method=RangeMethod.LIST,
                                prange=[3, 5, 8, 10, 15],
                                default=3),
        'min_impurity_decrease': 0.0,
        'init': None,
        'random_state': None,
        'max_features': ParamRange(method=RangeMethod.LIST,
                                   prange=['sqrt', None]),
        'verbose': 0,
        'max_leaf_nodes': None,
        'warm_start': False,
        'validation_fraction': 0.1,
        'n_iter_no_change': None,
        'tol': 0.0001,
        'ccp_alpha': 0.0,
    },
    # Regression Algorithms
    {
        'name': 'sklearn.ensemble.GradientBoostingRegressor',
        'tasks': [TaskType.REGRESSION],
        'data_types': [DataType.TABULAR],
        'tags': {'supports_random_seed': ['true']},
        # Hyperparams
        'loss': 'squared_error',
        'learning_rate': 0.1,
        'n_estimators': 100,
        'subsample': 1.0,
        'criterion': 'friedman_mse',
        'min_samples_split': ParamRange(method=RangeMethod.LIST,
                                        prange=[2, 5, 10],
                                        default=2),
        'min_samples_leaf': ParamRange(method=RangeMethod.LIST,
                                       prange=[1, 2, 5],
                                       default=1),
        'min_weight_fraction_leaf': 0.0,
        'max_depth': ParamRange(method=RangeMethod.LIST,
                                prange=[3, 5, 8, 10, 15],
                                default=3),
        'min_impurity_decrease': 0.0,
        'init': None,
        'max_features': ParamRange(method=RangeMethod.LIST,
                                   prange=['sqrt', None]),
        'alpha': 0.9,
        'verbose': 0,
        'max_leaf_nodes': None,
        'warm_start': False,
        'validation_fraction': 0.1,
        'n_iter_no_change': None,
        'tol': 1e-4,
        'ccp_alpha': 0.0,
    },
    {
        'name': 'sklearn.linear_model.LinearRegression',
        'tasks': [TaskType.REGRESSION],
        'data_types': [DataType.TABULAR],
        'tags': {'for_tests': ['true']},
        # Hyperparams
        'fit_intercept': True,
        'copy_X': True,
        'n_jobs': None,
        'positive': False,
    },
    {
        'name': 'sklearn.svm.SVR',
        'tasks': [TaskType.REGRESSION],
        'data_types': [DataType.TABULAR],
        # Hyperparams
        'kernel': 'rbf',
        'degree': 3,
        'gamma': 'scale',
        'coef0': 0.0,
        'tol': 1e-3,
        'C': ParamRange(method=RangeMethod.LIST,
                        prange=[0.01, 0.1, 1.0, 10.0, 100.0],
                        default=1.0),
        'epsilon': 0.1,
        'shrinking': True,
        'cache_size': 200,
        'verbose': False,
        'max_iter': -1,
    },
    {
        'name': 'sklearn.svm.LinearSVR',
        'tasks': [TaskType.REGRESSION],
        'data_types': [DataType.TABULAR],
        'tags': {'supports_random_seed': ['true']},
        # Hyperparams
        'epsilon': 0.0,
        'tol': 1e-4,
        'C': ParamRange(method=RangeMethod.LIST,
                        prange=[0.01, 0.1, 1.0, 10.0, 100.0],
                        default=1.0),
        'loss': 'epsilon_insensitive',
        'fit_intercept': True,
        'intercept_scaling': 1.0,
        'dual': True,
        'verbose': 0,
        'max_iter': 1000
    },
    {
        'name': 'sklearn.ensemble.RandomForestRegressor',
        'tasks': [TaskType.REGRESSION],
        'data_types': [DataType.TABULAR],
        'tags': {
            'for_tests': ['true'],
            'supports_random_seed': ['true']},
        # Hyperparams
        'n_estimators': 100,
        'criterion': 'squared_error',
        'max_depth': ParamRange(method=RangeMethod.LIST,
                                prange=[8, 10, 15, None]),
        'min_samples_split': ParamRange(method=RangeMethod.LIST,
                                        prange=[2, 5, 10],
                                        default=2),
        'min_samples_leaf': 1,
        'min_weight_fraction_leaf': 0.0,
        'max_features': 1.0,
        'max_leaf_nodes': None,
        'min_impurity_decrease': 0.0,
        'bootstrap': True,
        'oob_score': False,
        'n_jobs': None,
        'random_state': None,
        'verbose': 0,
        'warm_start': False,
        'ccp_alpha': 0.0,
        'max_samples': None,
    },
    {
        'name': 'sklearn.linear_model.Ridge',
        'tasks': [TaskType.REGRESSION],
        'data_types': [DataType.TABULAR],
        'tags': {'supports_random_seed': ['true']},
        # Hyperparams
        'alpha': ParamRange(method=RangeMethod.LIST,
                            prange=[0.001, 0.01, 0.1, 1.0, 5.0],
                            default=1.0),
        'fit_intercept': True,
        'copy_X': True,
        'max_iter': None,
        'tol': 0.0001,
        'solver': 'auto',
        'positive': False,
        'random_state': None,
    },
    {
        'name': 'sklearn.linear_model.Lasso',
        'tasks': [TaskType.REGRESSION],
        'data_types': [DataType.TABULAR],
        'tags': {'supports_random_seed': ['true']},
        # Hyperparams
        'alpha': ParamRange(method=RangeMethod.LIST,
                            prange=[0.001, 0.01, 0.1, 1.0, 5.0],
                            default=1.0),
        'fit_intercept': True,
        'precompute': False,
        'copy_X': True,
        'max_iter': 1000,
        'tol': 0.0001,
        'warm_start': False,
        'positive': False,
        'random_state': None,
        'selection': 'cyclic'
    },
    {
        'name': 'sklearn.linear_model.ElasticNet',
        'tasks': [TaskType.REGRESSION],
        'data_types': [DataType.TABULAR],
        'tags': {'supports_random_seed': ['true']},
        # Hyperparams
        'alpha': 1.0,
        'l1_ratio': 0.5,
        'fit_intercept': True,
        'precompute': False,
        'max_iter': 1000,
        'copy_X': True,
        'tol': 0.0001,
        'warm_start': False,
        'positive': False,
        'random_state': None,
        'selection': 'cyclic'
    },
    {
        'name': 'sklearn.linear_model.LassoCV',
        'tasks': [TaskType.REGRESSION],
        'data_types': [DataType.TABULAR],
        'tags': {'supports_random_seed': ['true']},
        # Hyperparams
        'eps': 0.001,
        'n_alphas': 100,
        'alphas': None,
        'fit_intercept': True,
        'precompute': 'auto',
        'max_iter': 1000,
        'tol': 0.0001,
        'copy_X': True,
        'cv': None,
        'verbose': False,
        'n_jobs': None,
        'positive': False,
        'random_state': None,
        'selection': 'cyclic'
    },
    {
        'name': 'sklearn.ensemble.AdaBoostRegressor',
        'tasks': [TaskType.REGRESSION],
        'data_types': [DataType.TABULAR],
        'tags': {'supports_random_seed': ['true']},
        # Hyperparams
        'estimator': None,
        'n_estimators': 50,
        'learning_rate': 1.0,
        'loss': 'linear',
        'random_state': None,
    },
    {
        'name': 'sklearn.tree.ExtraTreeRegressor',
        'tasks': [TaskType.REGRESSION],
        'data_types': [DataType.TABULAR],
        'tags': {'supports_random_seed': ['true']},
        # Hyperparams
        'criterion': 'squared_error',
        'splitter': 'random',
        'max_depth': ParamRange(method=RangeMethod.LIST, prange=[8, 10, 15, None]),
        'min_samples_split': ParamRange(method=RangeMethod.LIST, prange=[2, 5, 10], default=2),
        'min_samples_leaf': 1,
        'min_weight_fraction_leaf': 0.0,
        'max_features': 1.0,
        'random_state': None,
        'min_impurity_decrease': 0.0,
        'max_leaf_nodes': None,
        'ccp_alpha': 0.0
    },
    {
        'name': 'sklearn.ensemble.BaggingRegressor',
        'tasks': [TaskType.REGRESSION],
        'data_types': [DataType.TABULAR],
        'tags': {'supports_random_seed': ['true']},
        # Hyperparams
        'estimator': None,
        'n_estimators': 10,
        'max_samples': 1.0,
        'max_features': 1.0,
        'bootstrap': True,
        'bootstrap_features': False,
        'oob_score': False,
        'warm_start': False,
        'n_jobs': None,
        'random_state': None,
        'verbose': 0
    },
    {
        'name': 'sklearn.mixture.GaussianMixture',
        'tasks': [TaskType.DENSITY_ESTIMATION],
        'data_types': [DataType.TABULAR],
        'tags': {'supports_random_seed': ['true']},
        # Hyperparams
        'n_components': 1,
        'covariance_type': 'full',
        'tol': 0.001,
        'reg_covar': 1e-06,
        'max_iter': 100,
        'n_init': 1,
        'init_params': 'kmeans',
        'weights_init': None,
        'means_init': None,
        'precisions_init': None,
        'random_state': None,
        'warm_start': False,
        'verbose': 0,
        'verbose_interval': 10,
    },
    {
        'name': 'sklearn.tree.DecisionTreeClassifier',
        'tasks': [TaskType.BINARY_CLASSIFICATION, TaskType.MULTICLASS_CLASSIFICATION],
        'data_types': [DataType.TABULAR],
        'tags': {'supports_random_seed': ['true']},
        # Hyperparams
        'criterion': 'gini',
        'splitter': 'best',
        'max_depth': None,
        'min_samples_split': 2,
        'min_samples_leaf': 1,
        'min_weight_fraction_leaf': 0.0,
        'max_features': None,
        'random_state': None,
        'max_leaf_nodes': None,
        'min_impurity_decrease': 0.0,
        'class_weight': None,
        'ccp_alpha': 0.0,
        'monotonic_cst': None,
    },
]


def register(catalog: MemoryAlgorithmCatalog, **kwargs):
    '''Register all the algorithms in this file.'''
    for alg in ALGORITHMS:
        new_args = alg.copy()
        new_args.update(kwargs)
        catalog.register(SklearnAlgorithm(**new_args))
