import numpy as np
import pytest
from sklearn.cluster import MiniBatchKMeans
from sklearn.utils.estimator_checks import check_estimator

from nlp_playground.exceptions import BadValidation
from nlp_playground.lib.sklearn.estimators import ClusteringWithSupervision
from nlp_playground.lib.sklearn.estimators import DummyEstimator


def make_X_y():
    """Return a dummy X (training data) and y(classes)."""
    return (
        np.random.random(100).reshape(10, 10),
        np.arange(10) // 3
    )


def test__DummyEstimator_interface():
    check_estimator(DummyEstimator)


def test__ClusteringWithSupervision_interface():
    """Test that we are a legit sklearn estimator."""
    check_estimator(ClusteringWithSupervision)


def test__ClusteringWithSupervision_validation():
    train, classes = make_X_y()
    model = ClusteringWithSupervision()
    with pytest.raises(BadValidation):
        model.fit(train, classes)


def test__ClusteringWithSupervision_clusters():
    """
    Check that we change the number of clusters properly.

    We have a weird interface here where we sort of overload
    `n_clusters` but try to hide it.
    """
    train, classes = make_X_y()
    model = ClusteringWithSupervision(cluster_instance=MiniBatchKMeans())
    assert model.n_clusters is None
    assert model.get_params()['n_clusters'] is None
    assert model.cluster_instance.n_clusters == 8
    assert model._cluster_instance is None
    model.fit(train, classes)
    assert model.n_clusters is None
    assert model.get_params()['n_clusters'] is None
    assert model.cluster_instance.n_clusters == 8
    assert model._cluster_instance.n_clusters == 4
