# 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, 1, 54, 68), 'pytest.mark.parametrize', 'pytest.mark.parametrize', ({(54, 25, 54, 37): '"""func_tuple"""', (54, 39, 54, 67): 'supported_metrics_unweighted'}, {}), "('func_tuple', supported_metrics_unweighted)", False, 'import pytest\n'), ((71, 1, 71, 73), 'pytest.mark.parametrize', 'pytest.mark.parametrize', ({(71, 25, 71, 37): '"""func_tuple"""', (71, 39, 71, 72): 'supported_metrics_weighted_binary'}, {}), "('func_tuple', supported_metrics_weighted_binary)", False, 'import pytest\n'), ((86, 1, 86, 66), 'pytest.mark.parametrize', 'pytest.mark.parametrize', ({(86, 25, 86, 37): '"""func_tuple"""', (86, 39, 86, 65): 'supported_metrics_weighted'}, {}), "('func_tuple', supported_metrics_weighted)", False, 'import pytest\n'), ((104, 13, 104, 82), 'fairlearn.metrics.group_accuracy_score', 'metrics.group_accuracy_score', (), '', True, 'import fairlearn.metrics as metrics\n'), ((106, 23, 106, 64), 'sklearn.metrics.accuracy_score', 'skm.accuracy_score', ({(106, 42, 106, 48): 'Y_true', (106, 50, 106, 56): 'Y_pred', (106, 58, 106, 63): 'False'}, {}), '(Y_true, Y_pred, False)', True, 'import sklearn.metrics as skm\n'), ((116, 13, 116, 82), 'fairlearn.metrics.group_confusion_matrix', 'metrics.group_confusion_matrix', (), '', True, 'import fairlearn.metrics as metrics\n'), ((117, 23, 117, 74), 'sklearn.metrics.confusion_matrix', 'skm.confusion_matrix', (), '', True, 'import sklearn.metrics as skm\n'), ((119, 11, 119, 59), 'numpy.array_equal', 'np.array_equal', ({(119, 26, 119, 40): 'result.overall', (119, 42, 119, 58): 'expected_overall'}, {}), '(result.overall, expected_overall)', True, 'import numpy as np\n'), ((125, 13, 125, 96), 'fairlearn.metrics.group_precision_score', 'metrics.group_precision_score', (), '', True, 'import fairlearn.metrics as metrics\n'), ((126, 23, 126, 88), 'sklearn.metrics.precision_score', 'skm.precision_score', (), '', True, 'import sklearn.metrics as skm\n'), ((128, 11, 128, 59), 'numpy.array_equal', 'np.array_equal', ({(128, 26, 128, 40): 'result.overall', (128, 42, 128, 58): 'expected_overall'}, {}), '(result.overall, expected_overall)', True, 'import numpy as np\n'), ((132, 13, 132, 79), 'fairlearn.metrics.group_precision_score', 'metrics.group_precision_score', (), '', True, 'import fairlearn.metrics as metrics\n'), ((133, 23, 133, 71), 'sklearn.metrics.precision_score', 'skm.precision_score', (), '', True, 'import sklearn.metrics as skm\n'), ((135, 11, 135, 59), 'numpy.array_equal', 'np.array_equal', ({(135, 26, 135, 40): 'result.overall', (135, 42, 135, 58): 'expected_overall'}, {}), '(result.overall, expected_overall)', True, 'import numpy as np\n'), ((141, 13, 141, 93), 'fairlearn.metrics.group_recall_score', 'metrics.group_recall_score', (), '', True, 'import fairlearn.metrics as metrics\n'), ((142, 23, 142, 85), 'sklearn.metrics.recall_score', 'skm.recall_score', (), '', True, 'import sklearn.metrics as skm\n'), ((144, 11, 144, 59), 'numpy.array_equal', 'np.array_equal', ({(144, 26, 144, 40): 'result.overall', (144, 42, 144, 58): 'expected_overall'}, {}), '(result.overall, expected_overall)', True, 'import numpy as np\n'), ((148, 13, 148, 76), 'fairlearn.metrics.group_recall_score', 'metrics.group_recall_score', (), '', True, 'import fairlearn.metrics as metrics\n'), ((149, 23, 149, 68), 'sklearn.metrics.recall_score', 'skm.recall_score', (), '', True, 'import sklearn.metrics as skm\n'), ((151, 11, 151, 59), 'numpy.array_equal', 'np.array_equal', ({(151, 26, 151, 40): 'result.overall', (151, 42, 151, 58): 'expected_overall'}, {}), '(result.overall, expected_overall)', True, 'import numpy as np\n'), ((157, 13, 157, 83), 'fairlearn.metrics.group_roc_auc_score', 'metrics.group_roc_auc_score', (), '', True, 'import fairlearn.metrics as metrics\n'), ((158, 23, 158, 75), 'sklearn.metrics.roc_auc_score', 'skm.roc_auc_score', (), '', True, 'import sklearn.metrics as skm\n'), ((164, 13, 164, 77), 'fairlearn.metrics.group_roc_auc_score', 'metrics.group_roc_auc_score', (), '', True, 'import fairlearn.metrics as metrics\n'), ((165, 23, 165, 69), 'sklearn.metrics.roc_auc_score', 'skm.roc_auc_score', (), '', True, 'import sklearn.metrics as skm\n'), ((173, 13, 173, 81), 'fairlearn.metrics.group_zero_one_loss', 'metrics.group_zero_one_loss', (), '', True, 'import fairlearn.metrics as metrics\n'), ((175, 23, 175, 63), 'sklearn.metrics.zero_one_loss', 'skm.zero_one_loss', ({(175, 41, 175, 47): 'Y_true', (175, 49, 175, 55): 'Y_pred', (175, 57, 175, 62): 'False'}, {}), '(Y_true, Y_pred, False)', True, 'import sklearn.metrics as skm\n'), ((185, 13, 185, 89), 'fairlearn.metrics.group_mean_squared_error', 'metrics.group_mean_squared_error', (), '', True, 'import fairlearn.metrics as metrics\n'), ((187, 23, 187, 81), 'sklearn.metrics.mean_squared_error', 'skm.mean_squared_error', (), '', True, 'import sklearn.metrics as skm\n'), ((189, 11, 189, 59), 'numpy.array_equal', 'np.array_equal', ({(189, 26, 189, 40): 'result.overall', (189, 42, 189, 58): 'expected_overall'}, {}), '(result.overall, expected_overall)', True, 'import numpy as np\n'), ((197, 13, 197, 79), 'fairlearn.metrics.group_r2_score', 'metrics.group_r2_score', (), '', True, 'import fairlearn.metrics as metrics\n'), ((199, 23, 199, 71), 'sklearn.metrics.r2_score', 'skm.r2_score', (), '', True, 'import sklearn.metrics as skm\n'), ((201, 11, 201, 59), 'numpy.array_equal', 'np.array_equal', ({(201, 26, 201, 40): 'result.overall', (201, 42, 201, 58): 'expected_overall'}, {}), '(result.overall, expected_overall)', True, 'import numpy as np\n'), ((202, 24, 202, 41), 'numpy.unique', 'np.unique', ({(202, 34, 202, 40): 'groups'}, {}), '(groups)', True, 'import numpy as np\n'), ((213, 13, 213, 89), 'fairlearn.metrics.group_mean_squared_error', 'metrics.group_mean_squared_error', (), '', True, 'import fairlearn.metrics as metrics\n'), ((215, 23, 215, 81), 'sklearn.metrics.mean_squared_error', 'skm.mean_squared_error', (), '', True, 'import sklearn.metrics as skm\n'), ((217, 11, 217, 59), 'numpy.array_equal', 'np.array_equal', ({(217, 26, 217, 40): 'result.overall', (217, 42, 217, 58): 'expected_overall'}, {}), '(result.overall, expected_overall)', True, 'import numpy as np\n'), ((219, 24, 219, 41), 'numpy.unique', 'np.unique', ({(219, 34, 219, 40): 'groups'}, {}), '(groups)', True, 'import numpy as np\n'), ((66, 15, 66, 63), 'numpy.array_equal', 'np.array_equal', ({(66, 30, 66, 46): 'expected_overall', (66, 48, 66, 62): 'result.overall'}, {}), '(expected_overall, result.overall)', True, 'import numpy as np\n'), ((81, 15, 81, 63), 'numpy.array_equal', 'np.array_equal', ({(81, 30, 81, 46): 'expected_overall', (81, 48, 81, 62): 'result.overall'}, {}), '(expected_overall, result.overall)', True, 'import numpy as np\n'), ((96, 15, 96, 63), 'numpy.array_equal', 'np.array_equal', ({(96, 30, 96, 46): 'expected_overall', (96, 48, 96, 62): 'result.overall'}, {}), '(expected_overall, result.overall)', True, 'import numpy as np\n'), ((204, 19, 204, 79), 'sklearn.metrics.r2_score', 'skm.r2_score', (), '', True, 'import sklearn.metrics as skm\n'), ((205, 15, 205, 70), 'numpy.array_equal', 'np.array_equal', ({(205, 30, 205, 59): 'result.by_group[target_group]', (205, 61, 205, 69): 'expected'}, {}), '(result.by_group[target_group], expected)', True, 'import numpy as np\n'), ((211, 11, 211, 28), 'numpy.random.rand', 'np.random.rand', ({(211, 26, 211, 27): '(2)'}, {}), '(2)', True, 'import numpy as np\n'), ((212, 11, 212, 28), 'numpy.random.rand', 'np.random.rand', ({(212, 26, 212, 27): '(2)'}, {}), '(2)', True, 'import numpy as np\n'), ((226, 19, 226, 83), 'sklearn.metrics.mean_squared_error', 'skm.mean_squared_error', (), '', True, 'import sklearn.metrics as skm\n'), ((227, 15, 227, 70), 'numpy.array_equal', 'np.array_equal', ({(227, 30, 227, 59): 'result.by_group[target_group]', (227, 61, 227, 69): 'expected'}, {}), '(result.by_group[target_group], expected)', True, 'import numpy as np\n'), ((203, 15, 203, 33), 'numpy.asarray', 'np.asarray', ({(203, 26, 203, 32): 'groups'}, {}), '(groups)', True, 'import numpy as np\n')]