import numpy as np
import pytest
from sklearn.cluster import DBSCAN
from sklearn.ensemble import RandomForestClassifier
from sklearn.linear_model import Ridge, LogisticRegression

from sklego.common import flatten
from sklego.meta import SubjectiveClassifier
from tests.conftest import general_checks, classifier_checks, select_tests


@pytest.mark.parametrize(
    "test_fn",
    select_tests(
        flatten([general_checks, classifier_checks]),
        exclude=['check_sample_weights_invariance']
        # outliers train wont work because we have two thresholds
    )
)
def test_estimator_checks_classification(test_fn):
    if test_fn.__name__ == "check_classifiers_classes":
        prior = {
            "one": 0.1,
            "two": 0.1,
            "three": 0.1,
            -1: 0.1,
            1: 0.6,
        }  # nonsensical prior to make sklearn check pass
    else:
        prior = {0: 0.7, 1: 0.2, 2: 0.1}

    # Some of the sklearn checkers generate random y data with 3 classes, so prior needs to have these classes
    estimator = SubjectiveClassifier(LogisticRegression(), prior)
    test_fn(SubjectiveClassifier.__name__, estimator)


@pytest.mark.parametrize(
    "classes, prior, cfm, first_class_posterior",
    [
        ([0, 1], [0.8, 0.2], [[90, 10], [10, 90]], 0.973),  # numeric classes
        (["a", "b"], [0.8, 0.2], [[90, 10], [10, 90]], 0.973),  # char classes
        ([False, True], [0.8, 0.2], [[90, 10], [10, 90]], 0.973),  # bool classes
        (
            ["a", "b", "c"],
            [0.8, 0.1, 0.1],
            [[80, 10, 10], [10, 90, 0], [0, 0, 100]],
            0.985,
        ),  # n classes
        (
            [0, 1],
            [0.8, 0.2],
            [[100, 0], [0, 100]],
            1.0,
        ),  # "perfect" confusion matrix (no FP) -> prior is ignored
        (
            [0, 1],
            [0.2, 0.8],
            [[0, 100], [0, 100]],
            0.2,
        ),  # failure to predict class by inner estimator
        (
            [0, 1, 2],
            [0.1, 0.1, 0.8],
            [[0, 0, 100], [0, 0, 100], [0, 0, 100]],
            0.1,
        ),  # extremely biased, n classes
        (
            [0, 1, 2],
            [0.2, 0.1, 0.7],
            [[80, 0, 20], [0, 0, 100], [10, 0, 90]],
            0.696,
        ),  # biased, n classes
    ],
)
def test_posterior_computation(mocker, classes, prior, cfm, first_class_posterior):
    def mock_confusion_matrix(y, y_pred):
        return np.array(cfm)

    mocker.patch("sklego.meta.subjective_classifier.confusion_matrix", side_effect=mock_confusion_matrix)
    mock_estimator = mocker.Mock(RandomForestClassifier())
    mock_estimator.classes_ = np.array(classes)
    subjective_model = SubjectiveClassifier(mock_estimator, dict(zip(classes, prior)))
    subjective_model.fit(np.zeros((10, 10)), np.array([classes[0]] * 10))
    assert (
        pytest.approx(subjective_model.posterior_matrix_[0, 0], 0.001)
        == first_class_posterior
    )
    assert np.isclose(subjective_model.posterior_matrix_.sum(axis=0), 1).all()


@pytest.mark.parametrize(
    "prior, y",
    [
        (
            {"a": 0.8, "b": 0.2},
            ["a", "c"],
        ),  # class from train data not defined in prior
        ({"a": 0.8, "b": 0.2}, [0, 1]),  # different data types
    ],
)
def test_fit_y_data_inconsistent_with_prior_failure_conditions(prior, y):
    with pytest.raises(ValueError) as exc:
        SubjectiveClassifier(RandomForestClassifier(), prior).fit(
            np.zeros((len(y), 2)), np.array(y)
        )

    assert str(exc.value).startswith("Training data is inconsistent with prior")


def test_to_discrete():
    assert np.isclose(
        SubjectiveClassifier._to_discrete(
            np.array([[1, 0], [0.8, 0.2], [0.5, 0.5], [0.2, 0.8]])
        ),
        np.array([[1, 0], [1, 0], [1, 0], [0, 1]]),
    ).all()


@pytest.mark.parametrize(
    "weights,y_hats,expected_probas",
    [
        (
            [0.8, 0.2],
            [[1, 0], [0.5, 0.5], [0.8, 0.2]],
            [[1, 0], [0.8, 0.2], [0.94, 0.06]],
        ),
        (
            [0.5, 0.5],
            [[1, 0], [0.5, 0.5], [0.8, 0.2]],
            [[1, 0], [0.5, 0.5], [0.8, 0.2]],
        ),
        ([[0.8, 0.2], [0.5, 0.5]], [[1, 0], [0.8, 0.2]], [[1, 0], [0.8, 0.2]]),
    ],
)
def test_weighted_proba(weights, y_hats, expected_probas):
    assert np.isclose(
        SubjectiveClassifier._weighted_proba(np.array(weights), np.array(y_hats)),
        np.array(expected_probas),
        atol=1e-02,
    ).all()


@pytest.mark.parametrize(
    "evidence_type,expected_probas",
    [
        ("predict_proba", [[0.94, 0.06], [1, 0], [0.8, 0.2], [0.5, 0.5]]),
        ("confusion_matrix", [[0.97, 0.03], [0.97, 0.03], [0.97, 0.03], [0.47, 0.53]]),
        ("both", [[0.99, 0.01], [1, 0], [0.97, 0.03], [0.18, 0.82]]),
    ],
)
def test_predict_proba(mocker, evidence_type, expected_probas):
    def mock_confusion_matrix(y, y_pred):
        return np.array([[80, 20], [10, 90]])

    mocker.patch("sklego.meta.subjective_classifier.confusion_matrix", side_effect=mock_confusion_matrix)
    mock_inner_estimator = mocker.Mock(RandomForestClassifier)
    mock_inner_estimator.predict_proba.return_value = np.array(
        [[0.8, 0.2], [1, 0], [0.5, 0.5], [0.2, 0.8]]
    )
    mock_inner_estimator.classes_ = np.array([0, 1])
    subjective_model = SubjectiveClassifier(
        mock_inner_estimator, {0: 0.8, 1: 0.2}, evidence=evidence_type
    )
    subjective_model.fit(np.zeros((10, 10)), np.zeros(10))
    posterior_probabilities = subjective_model.predict_proba(np.zeros((4, 2)))
    assert posterior_probabilities.shape == (4, 2)
    assert np.isclose(
        posterior_probabilities, np.array(expected_probas), atol=0.01
    ).all()


@pytest.mark.parametrize(
    "inner_estimator, prior, evidence, expected_error_msg",
    [
        (DBSCAN(), {"a": 1}, "both", "Invalid inner estimator"),
        (Ridge(), {"a": 1}, "predict_proba", "Invalid inner estimator"),
        (
            RandomForestClassifier(),
            {"a": 0.8, "b": 0.1},
            "confusion_matrix",
            "Invalid prior",
        ),
        (
            RandomForestClassifier(),
            {"a": 0.8, "b": 0.2},
            "foo_evidence",
            "Invalid evidence",
        ),
    ],
)
def test_params_failure_conditions(
    inner_estimator, prior, evidence, expected_error_msg
):
    with pytest.raises(ValueError) as exc:
        SubjectiveClassifier(inner_estimator, prior, evidence).fit(
            np.zeros((2, 2)), np.zeros(2)
        )

    assert str(exc.value).startswith(expected_error_msg)
