# 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)

[((2279, 2346), 'pytest.mark.parametrize', 'pytest.mark.parametrize', (['"""func_tuple"""', 'supported_metrics_unweighted'], {}), "('func_tuple', supported_metrics_unweighted)\n", (2302, 2346), False, 'import pytest\n'), ((2896, 2968), 'pytest.mark.parametrize', 'pytest.mark.parametrize', (['"""func_tuple"""', 'supported_metrics_weighted_binary'], {}), "('func_tuple', supported_metrics_weighted_binary)\n", (2919, 2968), False, 'import pytest\n'), ((3442, 3507), 'pytest.mark.parametrize', 'pytest.mark.parametrize', (['"""func_tuple"""', 'supported_metrics_weighted'], {}), "('func_tuple', supported_metrics_weighted)\n", (3465, 3507), False, 'import pytest\n'), ((4169, 4238), 'fairlearn.metrics.group_accuracy_score', 'metrics.group_accuracy_score', (['Y_true', 'Y_pred', 'groups'], {'normalize': '(False)'}), '(Y_true, Y_pred, groups, normalize=False)\n', (4197, 4238), True, 'import fairlearn.metrics as metrics\n'), ((4263, 4304), 'sklearn.metrics.accuracy_score', 'skm.accuracy_score', (['Y_true', 'Y_pred', '(False)'], {}), '(Y_true, Y_pred, False)\n', (4281, 4304), True, 'import sklearn.metrics as skm\n'), ((4520, 4589), 'fairlearn.metrics.group_confusion_matrix', 'metrics.group_confusion_matrix', (['Y_true', 'Y_pred', 'groups'], {'labels': 'labels'}), '(Y_true, Y_pred, groups, labels=labels)\n', (4550, 4589), True, 'import fairlearn.metrics as metrics\n'), ((4613, 4664), 'sklearn.metrics.confusion_matrix', 'skm.confusion_matrix', (['Y_true', 'Y_pred'], {'labels': 'labels'}), '(Y_true, Y_pred, labels=labels)\n', (4633, 4664), True, 'import sklearn.metrics as skm\n'), ((4677, 4725), 'numpy.array_equal', 'np.array_equal', (['result.overall', 'expected_overall'], {}), '(result.overall, expected_overall)\n', (4691, 4725), True, 'import numpy as np\n'), ((4873, 4960), 'fairlearn.metrics.group_precision_score', 'metrics.group_precision_score', (['Y_true_ternary', 'Y_pred_ternary', 'group2'], {'average': 'None'}), '(Y_true_ternary, Y_pred_ternary, group2,\n    average=None)\n', (4902, 4960), True, 'import fairlearn.metrics as metrics\n'), ((4980, 5045), 'sklearn.metrics.precision_score', 'skm.precision_score', (['Y_true_ternary', 'Y_pred_ternary'], {'average': 'None'}), '(Y_true_ternary, Y_pred_ternary, average=None)\n', (4999, 5045), True, 'import sklearn.metrics as skm\n'), ((5058, 5106), 'numpy.array_equal', 'np.array_equal', (['result.overall', 'expected_overall'], {}), '(result.overall, expected_overall)\n', (5072, 5106), True, 'import numpy as np\n'), ((5166, 5232), 'fairlearn.metrics.group_precision_score', 'metrics.group_precision_score', (['Y_true', 'Y_pred', 'groups'], {'pos_label': '(0)'}), '(Y_true, Y_pred, groups, pos_label=0)\n', (5195, 5232), True, 'import fairlearn.metrics as metrics\n'), ((5256, 5304), 'sklearn.metrics.precision_score', 'skm.precision_score', (['Y_true', 'Y_pred'], {'pos_label': '(0)'}), '(Y_true, Y_pred, pos_label=0)\n', (5275, 5304), True, 'import sklearn.metrics as skm\n'), ((5317, 5365), 'numpy.array_equal', 'np.array_equal', (['result.overall', 'expected_overall'], {}), '(result.overall, expected_overall)\n', (5331, 5365), True, 'import numpy as np\n'), ((5510, 5595), 'fairlearn.metrics.group_recall_score', 'metrics.group_recall_score', (['Y_true_ternary', 'Y_pred_ternary', 'group2'], {'average': 'None'}), '(Y_true_ternary, Y_pred_ternary, group2, average=None\n    )\n', (5536, 5595), True, 'import fairlearn.metrics as metrics\n'), ((5614, 5676), 'sklearn.metrics.recall_score', 'skm.recall_score', (['Y_true_ternary', 'Y_pred_ternary'], {'average': 'None'}), '(Y_true_ternary, Y_pred_ternary, average=None)\n', (5630, 5676), True, 'import sklearn.metrics as skm\n'), ((5689, 5737), 'numpy.array_equal', 'np.array_equal', (['result.overall', 'expected_overall'], {}), '(result.overall, expected_overall)\n', (5703, 5737), True, 'import numpy as np\n'), ((5794, 5857), 'fairlearn.metrics.group_recall_score', 'metrics.group_recall_score', (['Y_true', 'Y_pred', 'groups'], {'pos_label': '(0)'}), '(Y_true, Y_pred, groups, pos_label=0)\n', (5820, 5857), True, 'import fairlearn.metrics as metrics\n'), ((5881, 5926), 'sklearn.metrics.recall_score', 'skm.recall_score', (['Y_true', 'Y_pred'], {'pos_label': '(0)'}), '(Y_true, Y_pred, pos_label=0)\n', (5897, 5926), True, 'import sklearn.metrics as skm\n'), ((5939, 5987), 'numpy.array_equal', 'np.array_equal', (['result.overall', 'expected_overall'], {}), '(result.overall, expected_overall)\n', (5953, 5987), True, 'import numpy as np\n'), ((6133, 6203), 'fairlearn.metrics.group_roc_auc_score', 'metrics.group_roc_auc_score', (['Y_true', 'Y_pred', 'groups'], {'average': '"""samples"""'}), "(Y_true, Y_pred, groups, average='samples')\n", (6160, 6203), True, 'import fairlearn.metrics as metrics\n'), ((6227, 6279), 'sklearn.metrics.roc_auc_score', 'skm.roc_auc_score', (['Y_true', 'Y_pred'], {'average': '"""samples"""'}), "(Y_true, Y_pred, average='samples')\n", (6244, 6279), True, 'import sklearn.metrics as skm\n'), ((6382, 6446), 'fairlearn.metrics.group_roc_auc_score', 'metrics.group_roc_auc_score', (['Y_true', 'Y_pred', 'groups'], {'max_fpr': '(0.5)'}), '(Y_true, Y_pred, groups, max_fpr=0.5)\n', (6409, 6446), True, 'import fairlearn.metrics as metrics\n'), ((6470, 6516), 'sklearn.metrics.roc_auc_score', 'skm.roc_auc_score', (['Y_true', 'Y_pred'], {'max_fpr': '(0.5)'}), '(Y_true, Y_pred, max_fpr=0.5)\n', (6487, 6516), True, 'import sklearn.metrics as skm\n'), ((6714, 6782), 'fairlearn.metrics.group_zero_one_loss', 'metrics.group_zero_one_loss', (['Y_true', 'Y_pred', 'groups'], {'normalize': '(False)'}), '(Y_true, Y_pred, groups, normalize=False)\n', (6741, 6782), True, 'import fairlearn.metrics as metrics\n'), ((6807, 6847), 'sklearn.metrics.zero_one_loss', 'skm.zero_one_loss', (['Y_true', 'Y_pred', '(False)'], {}), '(Y_true, Y_pred, False)\n', (6824, 6847), True, 'import sklearn.metrics as skm\n'), ((7153, 7229), 'fairlearn.metrics.group_mean_squared_error', 'metrics.group_mean_squared_error', (['y_t', 'y_p', 'groups'], {'multioutput': '"""raw_values"""'}), "(y_t, y_p, groups, multioutput='raw_values')\n", (7185, 7229), True, 'import fairlearn.metrics as metrics\n'), ((7254, 7312), 'sklearn.metrics.mean_squared_error', 'skm.mean_squared_error', (['y_t', 'y_p'], {'multioutput': '"""raw_values"""'}), "(y_t, y_p, multioutput='raw_values')\n", (7276, 7312), True, 'import sklearn.metrics as skm\n'), ((7325, 7373), 'numpy.array_equal', 'np.array_equal', (['result.overall', 'expected_overall'], {}), '(result.overall, expected_overall)\n', (7339, 7373), True, 'import numpy as np\n'), ((7607, 7673), 'fairlearn.metrics.group_r2_score', 'metrics.group_r2_score', (['y_t', 'y_p', 'groups'], {'multioutput': '"""raw_values"""'}), "(y_t, y_p, groups, multioutput='raw_values')\n", (7629, 7673), True, 'import fairlearn.metrics as metrics\n'), ((7698, 7746), 'sklearn.metrics.r2_score', 'skm.r2_score', (['y_t', 'y_p'], {'multioutput': '"""raw_values"""'}), "(y_t, y_p, multioutput='raw_values')\n", (7710, 7746), True, 'import sklearn.metrics as skm\n'), ((7759, 7807), 'numpy.array_equal', 'np.array_equal', (['result.overall', 'expected_overall'], {}), '(result.overall, expected_overall)\n', (7773, 7807), True, 'import numpy as np\n'), ((7832, 7849), 'numpy.unique', 'np.unique', (['groups'], {}), '(groups)\n', (7841, 7849), True, 'import numpy as np\n'), ((8318, 8394), 'fairlearn.metrics.group_mean_squared_error', 'metrics.group_mean_squared_error', (['y_t', 'y_p', 'groups'], {'multioutput': '"""raw_values"""'}), "(y_t, y_p, groups, multioutput='raw_values')\n", (8350, 8394), True, 'import fairlearn.metrics as metrics\n'), ((8419, 8477), 'sklearn.metrics.mean_squared_error', 'skm.mean_squared_error', (['y_t', 'y_p'], {'multioutput': '"""raw_values"""'}), "(y_t, y_p, multioutput='raw_values')\n", (8441, 8477), True, 'import sklearn.metrics as skm\n'), ((8490, 8538), 'numpy.array_equal', 'np.array_equal', (['result.overall', 'expected_overall'], {}), '(result.overall, expected_overall)\n', (8504, 8538), True, 'import numpy as np\n'), ((8564, 8581), 'numpy.unique', 'np.unique', (['groups'], {}), '(groups)\n', (8573, 8581), True, 'import numpy as np\n'), ((2784, 2832), 'numpy.array_equal', 'np.array_equal', (['expected_overall', 'result.overall'], {}), '(expected_overall, result.overall)\n', (2798, 2832), True, 'import numpy as np\n'), ((3330, 3378), 'numpy.array_equal', 'np.array_equal', (['expected_overall', 'result.overall'], {}), '(expected_overall, result.overall)\n', (3344, 3378), True, 'import numpy as np\n'), ((3909, 3957), 'numpy.array_equal', 'np.array_equal', (['expected_overall', 'result.overall'], {}), '(expected_overall, result.overall)\n', (3923, 3957), True, 'import numpy as np\n'), ((7920, 7980), 'sklearn.metrics.r2_score', 'skm.r2_score', (['y_t[mask]', 'y_p[mask]'], {'multioutput': '"""raw_values"""'}), "(y_t[mask], y_p[mask], multioutput='raw_values')\n", (7932, 7980), True, 'import sklearn.metrics as skm\n'), ((7996, 8051), 'numpy.array_equal', 'np.array_equal', (['result.by_group[target_group]', 'expected'], {}), '(result.by_group[target_group], expected)\n', (8010, 8051), True, 'import numpy as np\n'), ((8224, 8241), 'numpy.random.rand', 'np.random.rand', (['(2)'], {}), '(2)\n', (8238, 8241), True, 'import numpy as np\n'), ((8270, 8287), 'numpy.random.rand', 'np.random.rand', (['(2)'], {}), '(2)\n', (8284, 8287), True, 'import numpy as np\n'), ((8797, 8861), 'sklearn.metrics.mean_squared_error', 'skm.mean_squared_error', (['y_true', 'y_pred'], {'multioutput': '"""raw_values"""'}), "(y_true, y_pred, multioutput='raw_values')\n", (8819, 8861), True, 'import sklearn.metrics as skm\n'), ((8877, 8932), 'numpy.array_equal', 'np.array_equal', (['result.by_group[target_group]', 'expected'], {}), '(result.by_group[target_group], expected)\n', (8891, 8932), True, 'import numpy as np\n'), ((7866, 7884), 'numpy.asarray', 'np.asarray', (['groups'], {}), '(groups)\n', (7876, 7884), True, 'import numpy as np\n')]