import pytest
import numpy as np
from numpy.testing import assert_allclose
from numpy.testing import assert_array_equal

from sklearn.compose import make_column_transformer
from sklearn.datasets import make_classification
from sklearn.exceptions import NotFittedError
from sklearn.linear_model import LogisticRegression
from sklearn.pipeline import make_pipeline
from sklearn.preprocessing import StandardScaler
from sklearn.svm import SVC, SVR

from sklearn.metrics import confusion_matrix
from sklearn.metrics import plot_confusion_matrix
from sklearn.metrics import ConfusionMatrixDisplay


# TODO: Remove when https://github.com/numpy/numpy/issues/14397 is resolved
pytestmark = pytest.mark.filterwarnings(
    "ignore:In future, it will be an error for 'np.bool_':DeprecationWarning:"
    "matplotlib.*")


@pytest.fixture(scope="module")
def n_classes():
    return 5


@pytest.fixture(scope="module")
def data(n_classes):
    X, y = make_classification(n_samples=100, n_informative=5,
                               n_classes=n_classes, random_state=0)
    return X, y


@pytest.fixture(scope="module")
def fitted_clf(data):
    return SVC(kernel='linear', C=0.01).fit(*data)


@pytest.fixture(scope="module")
def y_pred(data, fitted_clf):
    X, _ = data
    return fitted_clf.predict(X)


def test_error_on_regressor(pyplot, data):
    X, y = data
    est = SVR().fit(X, y)

    msg = "plot_confusion_matrix only supports classifiers"
    with pytest.raises(ValueError, match=msg):
        plot_confusion_matrix(est, X, y)


def test_error_on_invalid_option(pyplot, fitted_clf, data):
    X, y = data
    msg = (r"normalize must be one of \{'true', 'pred', 'all', "
           r"None\}")

    with pytest.raises(ValueError, match=msg):
        plot_confusion_matrix(fitted_clf, X, y, normalize='invalid')


@pytest.mark.parametrize("with_labels", [True, False])
@pytest.mark.parametrize("with_display_labels", [True, False])
def test_plot_confusion_matrix_custom_labels(pyplot, data, y_pred, fitted_clf,
                                             n_classes, with_labels,
                                             with_display_labels):
    X, y = data
    ax = pyplot.gca()
    labels = [2, 1, 0, 3, 4] if with_labels else None
    display_labels = ['b', 'd', 'a', 'e', 'f'] if with_display_labels else None

    cm = confusion_matrix(y, y_pred, labels=labels)
    disp = plot_confusion_matrix(fitted_clf, X, y,
                                 ax=ax, display_labels=display_labels,
                                 labels=labels)

    assert_allclose(disp.confusion_matrix, cm)

    if with_display_labels:
        expected_display_labels = display_labels
    elif with_labels:
        expected_display_labels = labels
    else:
        expected_display_labels = list(range(n_classes))

    expected_display_labels_str = [str(name)
                                   for name in expected_display_labels]

    x_ticks = [tick.get_text() for tick in disp.ax_.get_xticklabels()]
    y_ticks = [tick.get_text() for tick in disp.ax_.get_yticklabels()]

    assert_array_equal(disp.display_labels, expected_display_labels)
    assert_array_equal(x_ticks, expected_display_labels_str)
    assert_array_equal(y_ticks, expected_display_labels_str)


@pytest.mark.parametrize("normalize", ['true', 'pred', 'all', None])
@pytest.mark.parametrize("include_values", [True, False])
def test_plot_confusion_matrix(pyplot, data, y_pred, n_classes, fitted_clf,
                               normalize, include_values):
    X, y = data
    ax = pyplot.gca()
    cmap = 'plasma'
    cm = confusion_matrix(y, y_pred)
    disp = plot_confusion_matrix(fitted_clf, X, y,
                                 normalize=normalize,
                                 cmap=cmap, ax=ax,
                                 include_values=include_values)

    assert disp.ax_ == ax

    if normalize == 'true':
        cm = cm / cm.sum(axis=1, keepdims=True)
    elif normalize == 'pred':
        cm = cm / cm.sum(axis=0, keepdims=True)
    elif normalize == 'all':
        cm = cm / cm.sum()

    assert_allclose(disp.confusion_matrix, cm)
    import matplotlib as mpl
    assert isinstance(disp.im_, mpl.image.AxesImage)
    assert disp.im_.get_cmap().name == cmap
    assert isinstance(disp.ax_, pyplot.Axes)
    assert isinstance(disp.figure_, pyplot.Figure)

    assert disp.ax_.get_ylabel() == "True label"
    assert disp.ax_.get_xlabel() == "Predicted label"

    x_ticks = [tick.get_text() for tick in disp.ax_.get_xticklabels()]
    y_ticks = [tick.get_text() for tick in disp.ax_.get_yticklabels()]

    expected_display_labels = list(range(n_classes))

    expected_display_labels_str = [str(name)
                                   for name in expected_display_labels]

    assert_array_equal(disp.display_labels, expected_display_labels)
    assert_array_equal(x_ticks, expected_display_labels_str)
    assert_array_equal(y_ticks, expected_display_labels_str)

    image_data = disp.im_.get_array().data
    assert_allclose(image_data, cm)

    if include_values:
        assert disp.text_.shape == (n_classes, n_classes)
        fmt = '.2g'
        expected_text = np.array([format(v, fmt) for v in cm.ravel(order="C")])
        text_text = np.array([
            t.get_text() for t in disp.text_.ravel(order="C")])
        assert_array_equal(expected_text, text_text)
    else:
        assert disp.text_ is None


def test_confusion_matrix_display(pyplot, data, fitted_clf, y_pred, n_classes):
    X, y = data

    cm = confusion_matrix(y, y_pred)
    disp = plot_confusion_matrix(fitted_clf, X, y, normalize=None,
                                 include_values=True, cmap='viridis',
                                 xticks_rotation=45.0)

    assert_allclose(disp.confusion_matrix, cm)
    assert disp.text_.shape == (n_classes, n_classes)

    rotations = [tick.get_rotation() for tick in disp.ax_.get_xticklabels()]
    assert_allclose(rotations, 45.0)

    image_data = disp.im_.get_array().data
    assert_allclose(image_data, cm)

    disp.plot(cmap='plasma')
    assert disp.im_.get_cmap().name == 'plasma'

    disp.plot(include_values=False)
    assert disp.text_ is None

    disp.plot(xticks_rotation=90.0)
    rotations = [tick.get_rotation() for tick in disp.ax_.get_xticklabels()]
    assert_allclose(rotations, 90.0)

    disp.plot(values_format='e')
    expected_text = np.array([format(v, 'e') for v in cm.ravel(order="C")])
    text_text = np.array([
        t.get_text() for t in disp.text_.ravel(order="C")])
    assert_array_equal(expected_text, text_text)


def test_confusion_matrix_contrast(pyplot):
    # make sure text color is appropriate depending on background

    cm = np.eye(2) / 2
    disp = ConfusionMatrixDisplay(cm, display_labels=[0, 1])

    disp.plot(cmap=pyplot.cm.gray)
    # diagonal text is black
    assert_allclose(disp.text_[0, 0].get_color(), [0.0, 0.0, 0.0, 1.0])
    assert_allclose(disp.text_[1, 1].get_color(), [0.0, 0.0, 0.0, 1.0])

    # off-diagonal text is white
    assert_allclose(disp.text_[0, 1].get_color(), [1.0, 1.0, 1.0, 1.0])
    assert_allclose(disp.text_[1, 0].get_color(), [1.0, 1.0, 1.0, 1.0])

    disp.plot(cmap=pyplot.cm.gray_r)
    # diagonal text is white
    assert_allclose(disp.text_[0, 1].get_color(), [0.0, 0.0, 0.0, 1.0])
    assert_allclose(disp.text_[1, 0].get_color(), [0.0, 0.0, 0.0, 1.0])

    # off-diagonal text is black
    assert_allclose(disp.text_[0, 0].get_color(), [1.0, 1.0, 1.0, 1.0])
    assert_allclose(disp.text_[1, 1].get_color(), [1.0, 1.0, 1.0, 1.0])

    # Regression test for #15920
    cm = np.array([[19, 34], [32, 58]])
    disp = ConfusionMatrixDisplay(cm, display_labels=[0, 1])

    disp.plot(cmap=pyplot.cm.Blues)
    min_color = pyplot.cm.Blues(0)
    max_color = pyplot.cm.Blues(255)
    assert_allclose(disp.text_[0, 0].get_color(), max_color)
    assert_allclose(disp.text_[0, 1].get_color(), max_color)
    assert_allclose(disp.text_[1, 0].get_color(), max_color)
    assert_allclose(disp.text_[1, 1].get_color(), min_color)


@pytest.mark.parametrize(
    "clf", [LogisticRegression(),
            make_pipeline(StandardScaler(), LogisticRegression()),
            make_pipeline(make_column_transformer((StandardScaler(), [0, 1])),
                          LogisticRegression())])
def test_confusion_matrix_pipeline(pyplot, clf, data, n_classes):
    X, y = data
    with pytest.raises(NotFittedError):
        plot_confusion_matrix(clf, X, y)
    clf.fit(X, y)
    y_pred = clf.predict(X)

    disp = plot_confusion_matrix(clf, X, y)
    cm = confusion_matrix(y, y_pred)

    assert_allclose(disp.confusion_matrix, cm)
    assert disp.text_.shape == (n_classes, n_classes)


@pytest.mark.parametrize("colorbar", [True, False])
def test_plot_confusion_matrix_colorbar(pyplot, data, fitted_clf, colorbar):
    X, y = data

    def _check_colorbar(disp, has_colorbar):
        if has_colorbar:
            assert disp.im_.colorbar is not None
            assert disp.im_.colorbar.__class__.__name__ == "Colorbar"
        else:
            assert disp.im_.colorbar is None
    disp = plot_confusion_matrix(fitted_clf, X, y, colorbar=colorbar)
    _check_colorbar(disp, colorbar)
    # attempt a plot with the opposite effect of colorbar
    disp.plot(colorbar=not colorbar)
    _check_colorbar(disp, not colorbar)


@pytest.mark.parametrize("values_format", ['e', 'n'])
def test_confusion_matrix_text_format(pyplot, data, y_pred, n_classes,
                                      fitted_clf, values_format):
    # Make sure plot text is formatted with 'values_format'.
    X, y = data
    cm = confusion_matrix(y, y_pred)
    disp = plot_confusion_matrix(fitted_clf, X, y,
                                 include_values=True,
                                 values_format=values_format)

    assert disp.text_.shape == (n_classes, n_classes)

    expected_text = np.array([format(v, values_format)
                              for v in cm.ravel()])
    text_text = np.array([
        t.get_text() for t in disp.text_.ravel()])
    assert_array_equal(expected_text, text_text)


def test_confusion_matrix_standard_format(pyplot):
    cm = np.array([[10000000, 0], [123456, 12345678]])
    plotted_text = ConfusionMatrixDisplay(
        cm, display_labels=[False, True]).plot().text_
    # Values should be shown as whole numbers 'd',
    # except the first number which should be shown as 1e+07 (longer length)
    # and the last number will be shown as 1.2e+07 (longer length)
    test = [t.get_text() for t in plotted_text.ravel()]
    assert test == ['1e+07', '0', '123456', '1.2e+07']

    cm = np.array([[0.1, 10], [100, 0.525]])
    plotted_text = ConfusionMatrixDisplay(
        cm, display_labels=[False, True]).plot().text_
    # Values should now formatted as '.2g', since there's a float in
    # Values are have two dec places max, (e.g 100 becomes 1e+02)
    test = [t.get_text() for t in plotted_text.ravel()]
    assert test == ['0.1', '10', '1e+02', '0.53']


@pytest.mark.parametrize("display_labels, expected_labels", [
    (None, ["0", "1"]),
    (["cat", "dog"], ["cat", "dog"]),
])
def test_default_labels(pyplot, display_labels, expected_labels):
    cm = np.array([[10, 0], [12, 120]])
    disp = ConfusionMatrixDisplay(cm, display_labels=display_labels).plot()

    x_ticks = [tick.get_text() for tick in disp.ax_.get_xticklabels()]
    y_ticks = [tick.get_text() for tick in disp.ax_.get_yticklabels()]

    assert_array_equal(x_ticks, expected_labels)
    assert_array_equal(y_ticks, expected_labels)


def test_error_on_a_dataset_with_unseen_labels(
    pyplot, fitted_clf, data, n_classes
):
    """Check that when labels=None, the unique values in `y_pred` and `y_true`
    will be used.
    Non-regression test for:
    https://github.com/scikit-learn/scikit-learn/pull/18405
    """
    X, y = data

    # create unseen labels in `y_true` not seen during fitting and not present
    # in 'fitted_clf.classes_'
    y = y + 1
    disp = plot_confusion_matrix(fitted_clf, X, y)

    display_labels = [tick.get_text() for tick in disp.ax_.get_xticklabels()]
    expected_labels = [str(i) for i in range(n_classes + 1)]
    assert_array_equal(expected_labels, display_labels)
