import numpy as np

import shap


def model(x):
    return np.array([np.linalg.norm(x)])


X = np.array([[3, 4], [5, 12], [7, 24]])
y = np.array([5, 13, 25])
explainer = np.array([[-1, 2], [-4, 2], [1, 2]])
masker = X


def test_update():
    """This is to test the update function within benchmark/framework"""
    sort_order = "positive"

    def score_function(true, pred):
        return np.mean(pred)

    perturbation = "keep"
    scores = {"name": "test", "metrics": list(), "values": dict()}

    shap.benchmark.update(model, X, y, explainer, masker, sort_order, score_function, perturbation, scores)

    metric = perturbation + " " + sort_order

    assert scores["metrics"][0] == metric
    assert len(scores["values"][metric]) == 3


def test_get_benchmark():
    """This is to test the get benchmark function within benchmark/framework"""
    metrics = {"sort_order": ["positive", "negative"], "perturbation": ["keep"]}
    scores = shap.benchmark.get_benchmark(model, X, y, explainer, masker, metrics)

    expected_metrics = ["keep positive", "keep negative"]

    assert set(expected_metrics) == set(scores["metrics"])
    assert len(scores["values"]) == 2


def test_get_metrics():
    """This is to test the get metrics function with respect to different selection method"""
    scores1 = {"name": "test1", "metrics": ["keep positive", "keep absolute"], "values": dict()}
    scores2 = {"name": "test2", "metrics": ["keep positive", "keep negative"], "values": dict()}
    benchmarks = {"test1": scores1, "test2": scores2}

    expected_metrics1 = set(["keep positive"])
    expected_metrics2 = set(["keep positive", "keep negative", "keep absolute"])

    assert set(shap.benchmark.get_metrics(benchmarks, lambda x, y: x.intersection(y))) == expected_metrics1
    assert set(shap.benchmark.get_metrics(benchmarks, lambda x, y: x.union(y))) == expected_metrics2