from .. import __version__ as sklearn_version

from ..decomposition.pca import PCA
from ..cluster.k_means_ import KMeans
from ..metrics import pairwise
from ..linear_model.base import LinearRegression
from ..linear_model.ridge import Ridge
from ..svm.classes import SVC

from .k_means import fit as kmeans_fit
from .k_means import predict as kmeans_predict
from .pca import _fit_full as pca_fit_full
from .pairwise import daal_pairwise_distances
from .linear import fit as linear_fit
from .linear import predict as linear_predict
from .ridge import fit as ridge_fit
from .ridge import predict as ridge_predict
from .svm import fit as svm_fit
from .svm import predict as svm_predict


mapping = {'PCA_full':  [[(PCA, '_fit_full', pca_fit_full), None]],
           'KMeans':    [[(KMeans, 'fit', kmeans_fit), None], [(KMeans, 'predict', kmeans_predict), None]],
           'distances': [[(pairwise, 'pairwise_distances', daal_pairwise_distances), None]],
           'linear':    [[(LinearRegression, 'fit', linear_fit), None],
                         [(LinearRegression, 'predict', linear_predict), None]],
           'ridge':     [[(Ridge, 'fit', ridge_fit), None],
                         [(Ridge, 'predict', ridge_predict), None]],
           'SVM':       [[(SVC, 'fit', svm_fit), None],
                         [(SVC, 'predict', svm_predict), None]], }


def do_patch(name):
    if name in mapping:
        for descriptor in mapping[name]:
            which, what, replacer = descriptor[0]
            if descriptor[1] is None:
                descriptor[1] = getattr(which, what)
            setattr(which, what, replacer)
    else:
        raise ValueError("Has no patch for: " + name)

def do_unpatch(name):
    if name in mapping:
        for descriptor in mapping[name]:
            which, what, replacer = descriptor[0]
            setattr(which, what, descriptor[1])
    else:
        raise ValueError("Has no patch for: " + name)

def enable(name=None):
    if sklearn_version != "0.19.1":
        raise NotImplementedError("daal4sklearn is for scikit-learn 0.19.1 only, found version {0}".format(sklearn_version))
    if name is not None:
        do_patch(name)
    else:
        for key in mapping:
            do_patch(key)

def disable(name=None):
    if name is not None:
        do_unpatch(name)
    else:
        for key in mapping:
            do_unpatch(key)
