# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License.

import pytest
import numpy as np
import sklearn.metrics as skm

import fairlearn.metrics as metrics

# ======================================================

a = "a"
b = "b"
c = "c"

Y_true = [0, 0, 0, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1]
Y_pred = [1, 0, 1, 0, 1, 0, 0, 1, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1]

Y_true_ternary = [a, b, c, c, c, b, b, b, c, c, a, a, a, a, a, b, c, c]
Y_pred_ternary = [b, c, c, c, b, b, b, b, b, c, a, a, c, a, a, b, c, c]

groups = [3, 4, 1, 0, 0, 0, 3, 2, 0, 1, 2, 3, 4, 0, 1, 2, 3, 4]
weight = [1, 2, 3, 1, 2, 3, 4, 2, 3, 3, 2, 1, 2, 3, 1, 2, 3, 4]
group2 = [0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1]

# =======================================================

# Define as a dictionary so that the actual name can be seen
# when pytest builds the tests

supported_metrics_weighted = [(skm.accuracy_score, metrics.group_accuracy_score),
                              (skm.confusion_matrix, metrics.group_confusion_matrix),
                              (skm.zero_one_loss, metrics.group_zero_one_loss)]

# The following only work with binary data when called with their default arguments
supported_metrics_weighted_binary = [(skm.precision_score, metrics.group_precision_score),
                                     (skm.recall_score, metrics.group_recall_score),
                                     (skm.roc_auc_score, metrics.group_roc_auc_score),
                                     (skm.mean_squared_error, metrics.group_mean_squared_error),
                                     (skm.r2_score, metrics.group_r2_score)]
supported_metrics_weighted_binary = supported_metrics_weighted_binary + supported_metrics_weighted


metrics_no_sample_weights = [(skm.max_error, metrics.group_max_error),
                             (skm.mean_absolute_error, metrics.group_mean_absolute_error),
                             (skm.mean_squared_log_error, metrics.group_mean_squared_log_error),
                             (skm.median_absolute_error, metrics.group_median_absolute_error)]

supported_metrics_unweighted = metrics_no_sample_weights + supported_metrics_weighted_binary

# =======================================================


@pytest.mark.parametrize("func_tuple", supported_metrics_unweighted)
def test_metric_unweighted(func_tuple):
    metric_func = func_tuple[0]
    group_metric_func = func_tuple[1]

    result = group_metric_func(Y_true, Y_pred, groups)
    # We don't really care about the numbers (sklearn is responsible)
    # We just want to make sure we got a result
    assert len(result.by_group) == 5

    expected_overall = metric_func(Y_true, Y_pred)
    if isinstance(expected_overall, np.ndarray):
        assert np.array_equal(expected_overall, result.overall)
    else:
        assert expected_overall == result.overall


@pytest.mark.parametrize("func_tuple", supported_metrics_weighted_binary)
def test_metric_weighted(func_tuple):
    metric_func = func_tuple[0]
    group_metric_func = func_tuple[1]

    result = group_metric_func(Y_true, Y_pred, groups, sample_weight=weight)
    assert len(result.by_group) == 5

    expected_overall = metric_func(Y_true, Y_pred, sample_weight=weight)
    if isinstance(expected_overall, np.ndarray):
        assert np.array_equal(expected_overall, result.overall)
    else:
        assert expected_overall == result.overall


@pytest.mark.parametrize("func_tuple", supported_metrics_weighted)
def test_metric_weighted_ternary(func_tuple):
    metric_func = func_tuple[0]
    group_metric_func = func_tuple[1]

    result = group_metric_func(Y_true_ternary, Y_pred_ternary, groups, sample_weight=weight)
    assert len(result.by_group) == 5
    expected_overall = metric_func(Y_true_ternary, Y_pred_ternary, sample_weight=weight)

    if isinstance(expected_overall, np.ndarray):
        assert np.array_equal(expected_overall, result.overall)
    else:
        assert expected_overall == result.overall


# ======================================================================================

def test_group_accuracy_score_unnormalized():
    result = metrics.group_accuracy_score(Y_true, Y_pred, groups, normalize=False)

    expected_overall = skm.accuracy_score(Y_true, Y_pred, False)

    assert result.overall == expected_overall


# ======================================================================================

def test_group_confusion_matrix_labels():
    labels = [0, 4]

    result = metrics.group_confusion_matrix(Y_true, Y_pred, groups, labels=labels)
    expected_overall = skm.confusion_matrix(Y_true, Y_pred, labels=labels)

    assert np.array_equal(result.overall, expected_overall)


# ======================================================================================

def test_group_precision_score_ternary():
    result = metrics.group_precision_score(Y_true_ternary, Y_pred_ternary, group2, average=None)
    expected_overall = skm.precision_score(Y_true_ternary, Y_pred_ternary, average=None)

    assert np.array_equal(result.overall, expected_overall)


def test_group_precision_score_pos_label():
    result = metrics.group_precision_score(Y_true, Y_pred, groups, pos_label=0)
    expected_overall = skm.precision_score(Y_true, Y_pred, pos_label=0)

    assert np.array_equal(result.overall, expected_overall)

# ======================================================================================


def test_group_recall_score_ternary():
    result = metrics.group_recall_score(Y_true_ternary, Y_pred_ternary, group2, average=None)
    expected_overall = skm.recall_score(Y_true_ternary, Y_pred_ternary, average=None)

    assert np.array_equal(result.overall, expected_overall)


def test_group_recall_score_pos_label():
    result = metrics.group_recall_score(Y_true, Y_pred, groups, pos_label=0)
    expected_overall = skm.recall_score(Y_true, Y_pred, pos_label=0)

    assert np.array_equal(result.overall, expected_overall)

# ======================================================================================


def test_group_roc_auc_score_average():
    result = metrics.group_roc_auc_score(Y_true, Y_pred, groups, average='samples')
    expected_overall = skm.roc_auc_score(Y_true, Y_pred, average='samples')

    assert expected_overall == result.overall


def test_group_roc_auc_score_max_fpr():
    result = metrics.group_roc_auc_score(Y_true, Y_pred, groups, max_fpr=0.5)
    expected_overall = skm.roc_auc_score(Y_true, Y_pred, max_fpr=0.5)

    assert expected_overall == result.overall

# ======================================================================================


def test_group_zero_one_loss_unnormalized():
    result = metrics.group_zero_one_loss(Y_true, Y_pred, groups, normalize=False)

    expected_overall = skm.zero_one_loss(Y_true, Y_pred, False)

    assert result.overall == expected_overall

# =============================================================================================


def test_group_mean_squared_error_multioutput_single_ndarray():
    y_t = np.random.rand(len(groups), 2)
    y_p = np.random.rand(len(groups), 2)
    result = metrics.group_mean_squared_error(y_t, y_p, groups, multioutput='raw_values')

    expected_overall = skm.mean_squared_error(y_t, y_p, multioutput='raw_values')

    assert np.array_equal(result.overall, expected_overall)

# =============================================================================================


def test_group_r2_score_multioutput():
    y_t = np.random.rand(len(groups), 2)
    y_p = np.random.rand(len(groups), 2)
    result = metrics.group_r2_score(y_t, y_p, groups, multioutput='raw_values')

    expected_overall = skm.r2_score(y_t, y_p, multioutput='raw_values')

    assert np.array_equal(result.overall, expected_overall)
    for target_group in np.unique(groups):
        mask = np.asarray(groups) == target_group
        expected = skm.r2_score(y_t[mask], y_p[mask], multioutput='raw_values')
        assert np.array_equal(result.by_group[target_group], expected)

# =============================================================================================


def test_group_mean_squared_error_multioutput_list_ndarray():
    y_t = [np.random.rand(2) for x in groups]
    y_p = [np.random.rand(2) for x in groups]
    result = metrics.group_mean_squared_error(y_t, y_p, groups, multioutput='raw_values')

    expected_overall = skm.mean_squared_error(y_t, y_p, multioutput='raw_values')

    assert np.array_equal(result.overall, expected_overall)

    for target_group in np.unique(groups):
        y_true = []
        y_pred = []
        for i in range(len(groups)):
            if groups[i] == target_group:
                y_true.append(y_t[i])
                y_pred.append(y_p[i])
        expected = skm.mean_squared_error(y_true, y_pred, multioutput='raw_values')
        assert np.array_equal(result.by_group[target_group], expected)

[(54, 'pytest.mark.parametrize', 'pytest.mark.parametrize', False, 'import pytest\n'), (71, 'pytest.mark.parametrize', 'pytest.mark.parametrize', False, 'import pytest\n'), (86, 'pytest.mark.parametrize', 'pytest.mark.parametrize', False, 'import pytest\n'), (104, 'fairlearn.metrics.group_accuracy_score', 'metrics.group_accuracy_score', True, 'import fairlearn.metrics as metrics\n'), (106, 'sklearn.metrics.accuracy_score', 'skm.accuracy_score', True, 'import sklearn.metrics as skm\n'), (116, 'fairlearn.metrics.group_confusion_matrix', 'metrics.group_confusion_matrix', True, 'import fairlearn.metrics as metrics\n'), (117, 'sklearn.metrics.confusion_matrix', 'skm.confusion_matrix', True, 'import sklearn.metrics as skm\n'), (119, 'numpy.array_equal', 'np.array_equal', True, 'import numpy as np\n'), (125, 'fairlearn.metrics.group_precision_score', 'metrics.group_precision_score', True, 'import fairlearn.metrics as metrics\n'), (126, 'sklearn.metrics.precision_score', 'skm.precision_score', True, 'import sklearn.metrics as skm\n'), (128, 'numpy.array_equal', 'np.array_equal', True, 'import numpy as np\n'), (132, 'fairlearn.metrics.group_precision_score', 'metrics.group_precision_score', True, 'import fairlearn.metrics as metrics\n'), (133, 'sklearn.metrics.precision_score', 'skm.precision_score', True, 'import sklearn.metrics as skm\n'), (135, 'numpy.array_equal', 'np.array_equal', True, 'import numpy as np\n'), (141, 'fairlearn.metrics.group_recall_score', 'metrics.group_recall_score', True, 'import fairlearn.metrics as metrics\n'), (142, 'sklearn.metrics.recall_score', 'skm.recall_score', True, 'import sklearn.metrics as skm\n'), (144, 'numpy.array_equal', 'np.array_equal', True, 'import numpy as np\n'), (148, 'fairlearn.metrics.group_recall_score', 'metrics.group_recall_score', True, 'import fairlearn.metrics as metrics\n'), (149, 'sklearn.metrics.recall_score', 'skm.recall_score', True, 'import sklearn.metrics as skm\n'), (151, 'numpy.array_equal', 'np.array_equal', True, 'import numpy as np\n'), (157, 'fairlearn.metrics.group_roc_auc_score', 'metrics.group_roc_auc_score', True, 'import fairlearn.metrics as metrics\n'), (158, 'sklearn.metrics.roc_auc_score', 'skm.roc_auc_score', True, 'import sklearn.metrics as skm\n'), (164, 'fairlearn.metrics.group_roc_auc_score', 'metrics.group_roc_auc_score', True, 'import fairlearn.metrics as metrics\n'), (165, 'sklearn.metrics.roc_auc_score', 'skm.roc_auc_score', True, 'import sklearn.metrics as skm\n'), (173, 'fairlearn.metrics.group_zero_one_loss', 'metrics.group_zero_one_loss', True, 'import fairlearn.metrics as metrics\n'), (175, 'sklearn.metrics.zero_one_loss', 'skm.zero_one_loss', True, 'import sklearn.metrics as skm\n'), (185, 'fairlearn.metrics.group_mean_squared_error', 'metrics.group_mean_squared_error', True, 'import fairlearn.metrics as metrics\n'), (187, 'sklearn.metrics.mean_squared_error', 'skm.mean_squared_error', True, 'import sklearn.metrics as skm\n'), (189, 'numpy.array_equal', 'np.array_equal', True, 'import numpy as np\n'), (197, 'fairlearn.metrics.group_r2_score', 'metrics.group_r2_score', True, 'import fairlearn.metrics as metrics\n'), (199, 'sklearn.metrics.r2_score', 'skm.r2_score', True, 'import sklearn.metrics as skm\n'), (201, 'numpy.array_equal', 'np.array_equal', True, 'import numpy as np\n'), (202, 'numpy.unique', 'np.unique', True, 'import numpy as np\n'), (213, 'fairlearn.metrics.group_mean_squared_error', 'metrics.group_mean_squared_error', True, 'import fairlearn.metrics as metrics\n'), (215, 'sklearn.metrics.mean_squared_error', 'skm.mean_squared_error', True, 'import sklearn.metrics as skm\n'), (217, 'numpy.array_equal', 'np.array_equal', True, 'import numpy as np\n'), (219, 'numpy.unique', 'np.unique', True, 'import numpy as np\n'), (66, 'numpy.array_equal', 'np.array_equal', True, 'import numpy as np\n'), (81, 'numpy.array_equal', 'np.array_equal', True, 'import numpy as np\n'), (96, 'numpy.array_equal', 'np.array_equal', True, 'import numpy as np\n'), (204, 'sklearn.metrics.r2_score', 'skm.r2_score', True, 'import sklearn.metrics as skm\n'), (205, 'numpy.array_equal', 'np.array_equal', True, 'import numpy as np\n'), (211, 'numpy.random.rand', 'np.random.rand', True, 'import numpy as np\n'), (212, 'numpy.random.rand', 'np.random.rand', True, 'import numpy as np\n'), (226, 'sklearn.metrics.mean_squared_error', 'skm.mean_squared_error', True, 'import sklearn.metrics as skm\n'), (227, 'numpy.array_equal', 'np.array_equal', True, 'import numpy as np\n'), (203, 'numpy.asarray', 'np.asarray', True, 'import numpy as np\n')]