import pytest

import numpy as np
from sklearn.utils.estimator_checks import check_estimator
from sklearn.metrics import confusion_matrix
from sklearn.model_selection import train_test_split

import legendary_potato.classifiers as classifiers


def classifier_iterator():
    """Return an iterator over classifier.
    """
    return (classifiers.SVDD,)


def two_class_generator(random_state=1576, dims=(1, 10)):
    """Generate classic multidimensional toy set.

    For each label, a set of centers are generated following normal
    distribution. A set of samples are generated following anormal distribution
    centered on this centers.

    The generated set is a 2-class balanced dataset.

    dims -- tuple indicating smallest and highest dimension tests
    """
    np.random.seed = random_state
    for dim in range(*dims):
        orig = np.zeros(dim)
        orig[0] = -1
        a_centers = np.random.normal(orig, 2, size=(3, dim))
        orig[0] = 1
        b_centers = np.random.normal(orig, 2, size=(3, dim))
        a_data = np.vstack(
            np.random.normal(center, 0.5, size=(5, dim))
            for center in a_centers
        )
        b_data = np.vstack(
            np.random.normal(center, 0.5, size=(5, dim))
            for center in b_centers
        )

        labels = [1 for _ in range(15)] + [-1 for _ in range(15)]
        data = np.hstack([np.vstack([a_data, b_data]), np.transpose([labels])])
        np.random.shuffle(data)
        yield data[:, 0:-1], data[:, -1]


@pytest.mark.parametrize('classifier', classifier_iterator())
@pytest.mark.parametrize('dataset', two_class_generator())
def test_oneclass(classifier, dataset):
    """Perform one class classification.
    """
    X, y = dataset
    X_train, X_test, y_train, y_test = train_test_split(X, y, shuffle=False)
    y_train = np.ones(X_train.shape[0])
    classif = classifier()
    classif.fit(X_train, y_train)
    classif.predict(X_test)


@pytest.mark.parametrize('classifier', classifier_iterator())
@pytest.mark.parametrize('dataset', two_class_generator())
def test_twoclasses(classifier, dataset):
    """Perform one class classification.
    """
    X, y = dataset
    X_train, X_test, y_train, y_test = train_test_split(X, y, shuffle=False)
    classif = classifier()
    classif.fit(X_train, y_train)
    y_pred = classif.predict(X_test)
    confusion_matrix(y_test, y_pred)


@pytest.mark.parametrize('classifier', classifier_iterator())
def test_sklearn_compatibility(classifier):
    """Check the compatibility.
    """
    check_estimator(classifier)

@@ -1,12 +1,43 @@
+from itertools import product%0A%0A
 import pytes
@@ -1535,16 +1535,888 @@
 , -1%5D%0A%0A%0A
+def multiclass_generator(random_state=1576, dims=(2, 3), nb_classes=(3, 5),%0A                         cl_sample=5):%0A    %22%22%22Generate classic multiclass toyset.%0A%0A    Works as the two_class_generator but alternativly with number and string%0A    labels.%0A    %22%22%22%0A    np.random.seed = random_state%0A    label_list = list('abcdefghijklmnopqrstuvwxyz')%0A    for dim, nb_cl in product(range(*dims), range(*nb_classes)):%0A        centers = np.random.normal(np.zeros(dim + 1), 5, size=(nb_cl, dim + 1))%0A        datas = np.vstack(%0A            np.random.normal(centers%5Bcl, :%5D, 0.5, size=(cl_sample, dim + 1))%0A            for cl in range(nb_cl)%0A        )%0A        datas%5B:, -1%5D = np.repeat(range(nb_cl), cl_sample)%0A        np.random.shuffle(datas)%0A        yield datas%5B:, 0:-1%5D, datas%5B:, -1%5D%0A        yield (datas%5B:, 0:-1%5D,%0A               np.array(%5Blabel_list%5Bint(i)%5D for i in datas%5B:, -1%5D%5D))%0A%0A%0A
 @pytest.
@@ -2957,32 +2957,32 @@
 ss_generator())%0A
-
 def test_twoclas
@@ -3285,24 +3285,470 @@
 , y_pred)%0A%0A%0A
+@pytest.mark.parametrize('classifier', classifier_iterator())%0A@pytest.mark.parametrize('dataset', multiclass_generator())%0Adef test_multiclass(classifier, dataset):%0A    %22%22%22Perform one class classification.%0A    %22%22%22%0A    X, y = dataset%0A    X_train, X_test, y_train, y_test = train_test_split(X, y, shuffle=False)%0A    classif = classifier()%0A    classif.fit(X_train, y_train)%0A    y_pred = classif.predict(X_test)%0A    confusion_matrix(y_test, y_pred)%0A%0A%0A
 @pytest.mark

Add multiclass test