#! /usr/bin/env python
# -*- coding: utf-8 -*-
'''
Created on 2022/8/18 22:59:51

@author: LiuKuan
@copyright: Apache License, Version 2.0
'''
from functools import partial
# ===============================================================================
''' 包装了 sklearn 所支持的数据集
'''
# ===============================================================================

dataset_list = [
    "make_biclusters",
    "make_blobs",
    "make_circles",
    "make_classification",
    "make_checkerboard",
    "make_friedman1",
    "make_friedman2",
    "make_friedman3",
    "make_gaussian_quantiles",
    "make_hastie_10_2",
    "make_low_rank_matrix",
    "make_moons",
    "make_multilabel_classification",
    "make_regression",
    "make_s_curve",
    "make_sparse_coded_signal",
    "make_sparse_spd_matrix",
    "make_sparse_uncorrelated",
    "make_spd_matrix",
    "make_swiss_roll",

    "load_boston",
    "load_diabetes",
    "load_digits",
    # "load_files",
    "load_iris",
    "load_breast_cancer",
    "load_linnerud",
    # "load_sample_image",
    "load_sample_images",
    # "load_svmlight_file",
    # "load_svmlight_files",
    "load_wine",

    "fetch_20newsgroups",
    "fetch_20newsgroups_vectorized",
    "fetch_lfw_pairs",
    "fetch_lfw_people",
    "fetch_olivetti_faces",
    "fetch_species_distributions",
    "fetch_california_housing",
    "fetch_covtype",
    "fetch_rcv1",
    "fetch_kddcup99",
    "fetch_openml",
]

def get_dataset(_name = 'make_blobs', **kws):
    name = _name.lower()
    ##### make
    if name == 'make_biclusters':
        from sklearn.datasets import make_biclusters
        dataset = partial(make_biclusters, shape = (2,2), n_clusters=4)
    elif name == 'make_blobs':
        from sklearn.datasets import make_blobs as dataset
    elif name == 'make_circles':
        from sklearn.datasets import make_circles as dataset
    elif name == 'make_classification':
        from sklearn.datasets import make_classification as dataset
    elif name == 'make_checkerboard':
        from sklearn.datasets import make_checkerboard
        dataset = partial(make_checkerboard, shape=(2, 2), n_clusters=4)
    elif name == 'make_friedman1':
        from sklearn.datasets import make_friedman1 as dataset
    elif name == 'make_friedman2':
        from sklearn.datasets import make_friedman2 as dataset
    elif name == 'make_friedman3':
        from sklearn.datasets import make_friedman3 as dataset
    elif name == 'make_gaussian_quantiles':
        from sklearn.datasets import make_gaussian_quantiles as dataset
    elif name == 'make_hastie_10_2':
        from sklearn.datasets import make_hastie_10_2 as dataset
    elif name == 'make_low_rank_matrix':
        from sklearn.datasets import make_low_rank_matrix as dataset
    elif name == 'make_moons':
        from sklearn.datasets import make_moons as dataset
    elif name == 'make_multilabel_classification':
        from sklearn.datasets import make_multilabel_classification as dataset
    elif name == 'make_regression':
        from sklearn.datasets import make_regression as dataset
    elif name == 'make_s_curve':
        from sklearn.datasets import make_s_curve as dataset
    elif name == 'make_sparse_coded_signal':
        from sklearn.datasets import make_sparse_coded_signal
        dataset = partial(make_sparse_coded_signal, n_samples = 100, n_components = 3, n_features = 2, n_nonzero_coefs = 0)
    elif name == 'make_sparse_spd_matrix':
        from sklearn.datasets import make_sparse_spd_matrix
        dataset = partial(make_sparse_spd_matrix, dim=1)
    elif name == 'make_sparse_uncorrelated':
        from sklearn.datasets import make_sparse_uncorrelated as dataset
    elif name == 'make_spd_matrix':
        from sklearn.datasets import make_spd_matrix
        dataset = partial(make_spd_matrix, n_dim=1)
    elif name == 'make_swiss_roll':
        from sklearn.datasets import make_swiss_roll as dataset

    ##### load
    elif name == 'load_boston':
        from sklearn.datasets import load_boston as dataset
    elif name == 'load_diabetes':
        from sklearn.datasets import load_diabetes as dataset
    elif name == 'load_digits':
        from sklearn.datasets import load_digits as dataset
    elif name == 'load_iris':
        from sklearn.datasets import load_iris as dataset
    elif name == 'load_breast_cancer':
        from sklearn.datasets import load_breast_cancer as dataset
    elif name == 'load_linnerud':
        from sklearn.datasets import load_linnerud as dataset
    elif name == 'load_sample_images':
        from sklearn.datasets import load_sample_images as dataset
    elif name == 'load_wine':
        from sklearn.datasets import load_wine as dataset

    ##### fetch
    elif name == 'fetch_20newsgroups':
        from sklearn.datasets import fetch_20newsgroups as dataset
    elif name == 'fetch_20newsgroups_vectorized':
        from sklearn.datasets import fetch_20newsgroups_vectorized as dataset
    elif name == 'fetch_lfw_pairs':
        from sklearn.datasets import fetch_lfw_pairs as dataset
    elif name == 'fetch_lfw_people':
        from sklearn.datasets import fetch_lfw_people as dataset
    elif name == 'fetch_olivetti_faces':
        from sklearn.datasets import fetch_olivetti_faces as dataset
    elif name == 'fetch_species_distributions':
        from sklearn.datasets import fetch_species_distributions as dataset
    elif name == 'fetch_california_housing':
        from sklearn.datasets import fetch_california_housing as dataset
    elif name == 'fetch_covtype':
        from sklearn.datasets import fetch_covtype as dataset
    elif name == 'fetch_rcv1':
        from sklearn.datasets import fetch_rcv1 as dataset
    elif name == 'fetch_kddcup99':
        from sklearn.datasets import fetch_kddcup99 as dataset
    elif name == 'fetch_openml':
        from sklearn.datasets import fetch_openml as dataset
    ##### 其他加载器
    elif name == 'load_files':
        from sklearn.datasets import load_files as dataset
    elif name == 'load_sample_image':
        from sklearn.datasets import load_sample_image as dataset
    elif name == 'load_svmlight_file':
        from sklearn.datasets import load_svmlight_file as dataset
    elif name == 'load_svmlight_files':
        from sklearn.datasets import load_svmlight_files as dataset
    else:
        raise RuntimeError("Cann't create {}".format(_name))

    return dataset(**kws)
