from collections import OrderedDict, defaultdict
from copy import deepcopy

import numpy as np
import pytest
from pytest_lazyfixture import lazy_fixture

from supervisor.api.defaults import DEFAULT_META_ANCHOR, DEFAULT_DATA_ANCHOR
from supervisor.exceptions import NotFittedError, PredictorCallError, PredictorReturnTypeError
from supervisor.explainers import AnchorTabular, DistributedAnchorTabular
from supervisor.explainers.tests.utils import predict_fcn


# TODO: Test DistributedAnchorBaseBeam separately
def uncollect_if_test_explainer(**kwargs):
    """
    This function is used to skip combinations of explainers
    and classifiers that do not make sense. This is achieved
    by using the hooks in conftest.py. Such functions should
    be passed to the @pytest.mark.uncollect_if decorator as
    the func argument. They should take the same inputs as the
    test function. Because this test function is parametrized
    with a lazy fixture, in this case the arguments name change
    (ie explainer can be both at_iris_explainer or at_adult_explainer),
    the **kwargs argument is used to collect all inputs.
    """

    # both explainer and classifier fixtures are parametrized using a LazyFixture
    # object that has a name attribute, representing the name of the dataset fixture
    clf_dataset_name = kwargs['rf_classifier'].name.split("_")[0]  # "adult_data" -> "adult"
    exp_dataset_name = kwargs['explainer'].name.split("_")[1]  # "at_adult_explainer" -> "adult"

    conditions = [
        clf_dataset_name != exp_dataset_name,
    ]

    return any(conditions)


@pytest.mark.uncollect_if(func=uncollect_if_test_explainer)
@pytest.mark.parametrize('n_explainer_runs', [5], ids='n_exp_runs={}'.format)
@pytest.mark.parametrize('at_defaults', [0.9, 0.95], ids='threshold={}'.format, indirect=True)
@pytest.mark.parametrize('rf_classifier',
                         [lazy_fixture('iris_data'), ],  # lazy_fixture('adult_data')],
                         indirect=True,
                         ids='clf=rf_{}'.format,
                         )
@pytest.mark.parametrize('explainer',
                         [lazy_fixture('at_iris_explainer'), ],  # lazy_fixture('at_adult_explainer')],
                         ids='exp={}'.format,
                         )
@pytest.mark.parametrize('test_instance_idx', [0], ids='test_instance_idx={}'.format)
def test_explainer(n_explainer_runs, at_defaults, rf_classifier, explainer, test_instance_idx, caplog):
    """
    Convergence test on Adult and Iris datasets.
    """

    # fixture returns a fitted AnchorTabular explainer
    X_test, explainer, predict_fn, predict_type = explainer
    if predict_type == 'proba':
        instance_label = np.argmax(predict_fn(X_test[test_instance_idx, :].reshape(1, -1)), axis=1)
    else:
        instance_label = predict_fn(X_test[test_instance_idx, :].reshape(1, -1))[0]

    explainer.instance_label = instance_label

    explain_defaults = at_defaults
    threshold = explain_defaults['desired_confidence']
    n_covered_ex = explain_defaults['n_covered_ex']

    run_precisions = []
    for _ in range(n_explainer_runs):
        explanation = explainer.explain(X_test[test_instance_idx], threshold=threshold, **explain_defaults)
        assert explainer.instance_label == instance_label
        if "Could not find" not in caplog.text:
            assert explanation.precision >= threshold
        assert explanation.coverage >= 0.001
        assert explanation.meta.keys() == DEFAULT_META_ANCHOR.keys()
        assert explanation.data.keys() == DEFAULT_DATA_ANCHOR.keys()
        run_precisions.append(explanation.precision)

    # check that 80% of runs returned a valid anchor
    # removed as the algorithm is too stochastic to pass for the fixed instance
    # assert ((np.asarray(run_precisions) > threshold).sum()) / n_explainer_runs >= 0.80

    sampler = explainer.samplers[0]
    assert sampler.instance_label == instance_label
    assert sampler.n_covered_ex == n_covered_ex


def uncollect_if_test_sampler(**kwargs):
    """
    Make sure only tests operating on the same dataset are run.
    """
    clf_dataset_name = kwargs['rf_classifier'].name.split("_")[0]  # "adult_data" -> "adult"
    exp_dataset_name = kwargs['explainer'].name.split("_")[1]  # "at_adult_explainer" -> "adult"
    dataset_name = kwargs['dataset'].name.split("_")[0]  # "adult_data" -> "adult"

    conditions = [
        len({clf_dataset_name, exp_dataset_name, dataset_name}) != 1,
    ]

    return any(conditions)


@pytest.mark.skip(reason='Not testing as performance of distributed anchors (within an instance) not clear.'
                         'Also, these tests fail intermittently with precision threshold not achieved.')
@pytest.mark.parametrize('ncpu', [2], ids='ncpu={}'.format)
@pytest.mark.parametrize('predict_type', ('proba', 'class'), ids='predict_type={}'.format)
@pytest.mark.parametrize('at_defaults', [0.9], ids='threshold={}'.format, indirect=True)
@pytest.mark.parametrize('rf_classifier',
                         [lazy_fixture('iris_data')],
                         indirect=True,
                         ids='clf=rf_{}'.format,
                         )
@pytest.mark.parametrize('test_instance_idx', [0], ids='test_instance_idx={}'.format)
def test_distributed_anchor_tabular(ncpu,
                                    predict_type,
                                    at_defaults,
                                    iris_data,
                                    rf_classifier,
                                    test_instance_idx,
                                    ):
    # TODO - if we add this test is back in, this conditional should be added as a @pytest.mark.skip
    # i.e. see test_kernel_distributed_execution in test_shap_wrappers.py.

    ray = pytest.importorskip('ray', reason="Distributed tests skipped as Ray not installed")

    # inputs
    params = at_defaults
    threshold = params['desired_confidence']
    n_covered_ex = params['n_covered_ex']  # number of covered examples to return when anchor applies
    batch_size = params['batch_size']  # number of samples to draw during sampling
    n_anchors_to_sample = 6  # for testing sampling function

    # prepare the classifier and explainer
    data = iris_data
    X_test, X_train, feature_names = data['X_test'], data['X_train'], data['metadata']['feature_names']
    clf, preprocessor = rf_classifier
    predictor = predict_fcn(predict_type, clf)
    explainer = DistributedAnchorTabular(predictor, feature_names, seed=0)
    explainer.fit(X_train, ncpu=ncpu)

    # select instance to be explained
    instance = X_test[test_instance_idx]
    if predict_type == 'proba':
        instance_label = np.argmax(predictor(instance.reshape(1, -1)), axis=1)
    else:
        instance_label = predictor(instance.reshape(1, -1))[0]

    # explain the instance and do basic checks on the lookups and instance labels used by samplers
    explanation = explainer.explain(instance, threshold=threshold, n_covered_ex=n_covered_ex)
    assert len(explainer.samplers) == ncpu
    actors = explainer.samplers
    for actor in actors:
        sampler = ray.get(actor._get_sampler.remote())
        ord_feats = sampler.ord_lookup.keys()
        cat_feats = sampler.cat_lookup.keys()
        enc_feats = sampler.enc2feat_idx.keys()
        assert (set(ord_feats | set(cat_feats))) == set(enc_feats)
        assert sampler.instance_label == instance_label
        assert sampler.n_covered_ex == n_covered_ex

    # check explanation
    assert explainer.instance_label == instance_label
    print(explanation.anchor)
    print(explanation.coverage)
    assert explanation.precision >= threshold
    assert explanation.coverage >= 0.05

    distrib_anchor_beam = explainer.mab
    assert len(distrib_anchor_beam.samplers) == ncpu

    # basic checks for DistributedAnchorBaseBeam
    anchor_features = list(enc_feats)
    anchor_max_len = len(anchor_features)
    assert distrib_anchor_beam.state['coverage_data'].shape[1] == anchor_max_len
    to_sample = []
    for _ in range(n_anchors_to_sample):
        anchor_len = np.random.randint(0, anchor_max_len)
        anchor = np.random.choice(anchor_features, anchor_len, replace=False)
        to_sample.append(tuple(anchor))
    to_sample = list(set(to_sample))
    current_state = deepcopy(distrib_anchor_beam.state)
    pos, total = distrib_anchor_beam.draw_samples(to_sample, batch_size)
    for p, t, anchor in zip(pos, total, to_sample):
        assert distrib_anchor_beam.state['t_nsamples'][anchor] == current_state['t_nsamples'][anchor] + t
        assert distrib_anchor_beam.state['t_positives'][anchor] == current_state['t_positives'][anchor] + p

    ray.shutdown()


@pytest.mark.uncollect_if(func=uncollect_if_test_sampler)
@pytest.mark.parametrize('test_instance_idx', [0], ids='test_instance_idx={}'.format)
@pytest.mark.parametrize('nb_samples', [100], ids='nb_samples={}'.format)
@pytest.mark.parametrize('anchors', [((2,), (10,), (11,), (7, 10, 11), (3, 11))], ids='anchors={}'.format)
@pytest.mark.parametrize('dataset',
                         [lazy_fixture('adult_data'), lazy_fixture('iris_data')],
                         ids='dataset={}'.format,
                         )
@pytest.mark.parametrize('rf_classifier',
                         [lazy_fixture('adult_data'), lazy_fixture('iris_data')],
                         indirect=True,
                         ids='clf=rf_{}'.format,
                         )
@pytest.mark.parametrize('explainer',
                         [lazy_fixture('at_iris_explainer'), lazy_fixture('at_adult_explainer')],
                         ids='exp={}'.format,
                         )
def test_sampler(test_instance_idx, anchors, nb_samples, dataset, rf_classifier, explainer):
    """
    Test sampler for datasets that contain continuos only or continuous and categorical variables.
    """

    # fixture returns a fitted AnchorTabular explainer
    X_test, explainer, predict_fn, predict_type = explainer
    test_dataset = dataset
    test_dataset_name = test_dataset['metadata']['name']

    # the `explain` call does some more sampler setup before using it so we need to repeat it here
    sampler = explainer.samplers[0]
    sampler.set_instance_label(X_test[test_instance_idx])
    # TODO: the test passes now, but we're not setting other things passed to `explain` like `n_covered_ex`
    #  or calling explainer._build_sampling_lookups. This suggests either the coupling between `sampler` and
    #  `explainer` is too strong or this test should be re-written in a different way.

    # test sampler setup is correct
    assert len(explainer.samplers) == 1
    sampler = explainer.samplers[0]
    assert explainer.predictor(X_test[test_instance_idx].reshape(1, -1)).shape == (1,)
    assert sampler.train_data.shape == sampler.d_train_data.shape == test_dataset['X_train'].shape

    if test_dataset_name == 'adult':
        category_map = test_dataset['metadata']['category_map']
        assert len(sampler.categorical_features) == len(category_map.keys())
        assert len(sampler.numerical_features) == X_test.shape[1] - len(category_map.keys())
    elif test_dataset_name == 'iris':
        assert (np.unique(sampler.d_train_data) == np.array([0., 1., 2., 3.])).all()
        assert not sampler.categorical_features
        assert len(sampler.numerical_features) == X_test.shape[1]

    if predict_type == 'proba':
        instance_label = np.argmax(predict_fn(X_test[test_instance_idx, :].reshape(1, -1)), axis=1)
    else:
        instance_label = predict_fn(X_test[test_instance_idx, :].reshape(1, -1))[0]

    explainer.instance_label = instance_label

    # test sampling function end2end
    train_data = sampler.train_data
    train_data_mean = np.mean(train_data, axis=0)[sampler.numerical_features]
    train_data_3std = 3 * np.std(train_data, axis=0)[sampler.numerical_features]
    sampler.build_lookups(X_test[test_instance_idx, :])
    n_covered_ex = sampler.n_covered_ex

    for anchor in anchors:
        cov_true, cov_false, labels, data, coverage, anchor_pos = sampler((0, anchor), nb_samples)
        assert cov_true.shape[0] <= n_covered_ex
        assert cov_false.shape[0] <= n_covered_ex
        assert len(labels) == nb_samples
        assert len(sampler.enc2feat_idx) == data.shape[1]
        assert coverage != -1

        # test lookups dictionary used for sampling
        ord_feats = sampler.ord_lookup.keys()
        cat_feats = sampler.cat_lookup.keys()
        enc_feats = sampler.enc2feat_idx.keys()
        assert (set(ord_feats | set(cat_feats))) == set(enc_feats)

        # now test perturbation method ...
        # find out which bins can be sampled for discretized vars and perturb some random data
        allowed_bins, allowed_rows, unk_feat_vals = sampler.get_features_index(anchor)
        raw_data, disc_data, coverage = sampler.perturbation(anchor, nb_samples)

        assert not unk_feat_vals
        assert coverage != -1
        assert raw_data.shape[0] == disc_data.shape[0] == nb_samples

        # find the indices of the features in the anchor
        uniq_feat_ids = list(OrderedDict.fromkeys([sampler.enc2feat_idx[enc_idx] for enc_idx in anchor]))
        uniq_feat_ids = [feat for feat in uniq_feat_ids if feat not in [f for f, _, _ in unk_feat_vals]]

        # find the mapping of feature ids to encoded feature ids
        feat2enc_idx = defaultdict(list)
        for enc_feat_idx, orig_feat_id in sampler.enc2feat_idx.items():
            feat2enc_idx[orig_feat_id].append(enc_feat_idx)

        # find the expected values the sampled features should have
        expected_values = []
        for feat_id in uniq_feat_ids:
            if feat_id in sampler.categorical_features:
                enc_idx = feat2enc_idx[feat_id][0]
                val = sampler.cat_lookup[enc_idx]
                expected_values.append({val})
            else:
                expected_values.append(allowed_bins[feat_id])

    # check that we have replaced features properly with values from the same bin
    # or with the correct value
    for bins, feat_id in zip(expected_values, uniq_feat_ids):
        sampled_bins_uniq = set(np.unique(disc_data[:, feat_id]))

        # check that we have replaced features properly with values from the same bin
        assert bins - sampled_bins_uniq == set()
        assert sampled_bins_uniq - bins == set()
        raw_data_mean = np.mean(raw_data, axis=0)[sampler.numerical_features]

        # check features sampled are in a sensible range for numerical features
        assert (train_data_mean + train_data_3std - raw_data_mean > 0).all()
        assert (train_data_mean - train_data_3std - raw_data_mean < 0).all()


def bad_predictor(x: np.ndarray) -> list:
    """
    A dummy predictor emulating the following:
     - Expecting an array of certain dimension
     - Returning an incorrect type
     This is used below to test custom exception functionality.
    """
    if x.shape[1] != 3:
        raise ValueError
    return list(x)


def good_predictor(x: np.ndarray) -> np.ndarray:
    """
    A dummy predictor returning a vector of random binary target labels.
    """
    return np.random.randint(low=0, high=2, size=x.shape[0])


def test_anchor_tabular_fails_init_bad_feature_names_predictor_call():
    """
    In this test `feature_names` is misspecified leading to an exception calling the `predictor`.
    """
    with pytest.raises(PredictorCallError):
        explainer = AnchorTabular(bad_predictor, feature_names=['f1', 'f2'])  # noqa: F841


def test_anchor_tabular_fails_bad_predictor_return_type():
    """
    In this test `feature_names` is specified correctly, but the predictor returns the wrong type.
    """
    with pytest.raises(PredictorReturnTypeError):
        explainer = AnchorTabular(bad_predictor, feature_names=['f1', 'f2', 'f3'])  # noqa: F841


def test_anchor_tabular_explain_fails_not_fitted():
    explainer = AnchorTabular(good_predictor, feature_names=['f1', 'f2'])
    with pytest.raises(NotFittedError) as err:
        explainer.explain(np.ones(2))
    expected_msg = "This AnchorTabular instance is not fitted yet. Call 'fit' with appropriate arguments first."
    assert str(err.value) == expected_msg
