"""Unit tests for directrs.models (DirectRSClassifier and DirectRSRegressor)."""

import numpy as np
import pytest
from sklearn.datasets import load_iris, load_diabetes

from directrs.models import DirectRSClassifier, DirectRSRegressor


# ---------------------------------------------------------------------------
# Fixtures
# ---------------------------------------------------------------------------

@pytest.fixture(scope="module")
def iris_data():
    return load_iris(return_X_y=True)


@pytest.fixture(scope="module")
def diabetes_data():
    return load_diabetes(return_X_y=True)


@pytest.fixture(scope="module")
def fitted_classifier(iris_data):
    X, y = iris_data
    clf = DirectRSClassifier(max_depth=4, random_state=0)
    clf.fit(X, y)
    return clf, X, y


@pytest.fixture(scope="module")
def fitted_regressor(diabetes_data):
    X, y = diabetes_data
    reg = DirectRSRegressor(max_depth=4, random_state=0)
    reg.fit(X, y)
    return reg, X, y


# ---------------------------------------------------------------------------
# DirectRSClassifier
# ---------------------------------------------------------------------------

class TestDirectRSClassifier:
    def test_fit_runs(self, iris_data):
        X, y = iris_data
        clf = DirectRSClassifier(max_depth=3, random_state=0)
        clf.fit(X, y)  # must not raise

    def test_predict_shape(self, fitted_classifier):
        clf, X, y = fitted_classifier
        preds = clf.predict(X)
        assert preds.shape == y.shape

    def test_predict_classes_valid(self, fitted_classifier):
        clf, X, y = fitted_classifier
        preds = clf.predict(X)
        assert set(preds).issubset(set(clf.classes_))

    def test_predict_proba_shape(self, fitted_classifier):
        clf, X, y = fitted_classifier
        proba = clf.predict_proba(X)
        assert proba.shape == (len(X), len(clf.classes_))

    def test_predict_proba_sums_to_one(self, fitted_classifier):
        clf, X, _ = fitted_classifier
        proba = clf.predict_proba(X)
        np.testing.assert_allclose(proba.sum(axis=1), 1.0, atol=1e-6)

    def test_predict_proba_nonnegative(self, fitted_classifier):
        clf, X, _ = fitted_classifier
        proba = clf.predict_proba(X)
        assert np.all(proba >= 0.0)

    def test_score_above_threshold(self, fitted_classifier):
        clf, X, y = fitted_classifier
        score = clf.score(X, y)
        assert score > 0.9, f"Expected accuracy > 0.9, got {score:.3f}"

    def test_exposes_G_matrix(self, fitted_classifier):
        clf, _, _ = fitted_classifier
        assert hasattr(clf, "G_")
        assert clf.G_.shape == (clf.n_features_in_, clf.n_features_in_)

    def test_exposes_S_prime(self, fitted_classifier):
        clf, _, _ = fitted_classifier
        assert hasattr(clf, "S_prime_")
        assert clf.S_prime_.shape == clf.G_.shape

    def test_exposes_leaf_indices(self, fitted_classifier):
        clf, _, _ = fitted_classifier
        assert hasattr(clf, "leaf_indices_")
        assert len(clf.leaf_indices_) > 0

    def test_feature_importances_keys(self, fitted_classifier):
        clf, _, _ = fitted_classifier
        fi = clf.feature_importances_
        assert "tree" in fi
        assert "geometry" in fi

    def test_feature_importances_shapes(self, fitted_classifier):
        clf, X, _ = fitted_classifier
        fi = clf.feature_importances_
        assert fi["tree"].shape == (X.shape[1],)
        assert fi["geometry"].shape == (X.shape[1],)

    def test_geometry_importance_sums_to_one(self, fitted_classifier):
        clf, _, _ = fitted_classifier
        geom = clf.feature_importances_["geometry"]
        np.testing.assert_allclose(geom.sum(), 1.0, atol=1e-10)

    def test_different_depths_work(self, iris_data):
        X, y = iris_data
        for depth in [1, 2, 3, 5]:
            clf = DirectRSClassifier(max_depth=depth, random_state=0)
            clf.fit(X, y)
            preds = clf.predict(X)
            assert len(preds) == len(y)

    def test_get_params_set_params(self):
        clf = DirectRSClassifier(max_depth=5, C=2.0)
        params = clf.get_params()
        assert params["max_depth"] == 5
        assert params["C"] == 2.0
        clf.set_params(max_depth=3)
        assert clf.max_depth == 3


# ---------------------------------------------------------------------------
# DirectRSRegressor
# ---------------------------------------------------------------------------

class TestDirectRSRegressor:
    def test_fit_runs(self, diabetes_data):
        X, y = diabetes_data
        reg = DirectRSRegressor(max_depth=3, random_state=0)
        reg.fit(X, y)  # must not raise

    def test_predict_shape(self, fitted_regressor):
        reg, X, y = fitted_regressor
        preds = reg.predict(X)
        assert preds.shape == y.shape

    def test_predict_is_float(self, fitted_regressor):
        reg, X, _ = fitted_regressor
        preds = reg.predict(X)
        assert preds.dtype.kind == "f"

    def test_r2_above_threshold(self, fitted_regressor):
        reg, X, y = fitted_regressor
        r2 = reg.score(X, y)
        assert r2 > 0.3, f"Expected R² > 0.3, got {r2:.3f}"

    def test_exposes_G_matrix(self, fitted_regressor):
        reg, _, _ = fitted_regressor
        assert hasattr(reg, "G_")
        assert reg.G_.shape == (reg.n_features_in_, reg.n_features_in_)

    def test_exposes_S_prime(self, fitted_regressor):
        reg, _, _ = fitted_regressor
        assert hasattr(reg, "S_prime_")
        assert reg.S_prime_.shape == reg.G_.shape

    def test_exposes_leaf_indices(self, fitted_regressor):
        reg, _, _ = fitted_regressor
        assert hasattr(reg, "leaf_indices_")
        assert len(reg.leaf_indices_) > 0

    def test_feature_importances_keys(self, fitted_regressor):
        reg, _, _ = fitted_regressor
        fi = reg.feature_importances_
        assert "tree" in fi
        assert "geometry" in fi

    def test_feature_importances_shapes(self, fitted_regressor):
        reg, X, _ = fitted_regressor
        fi = reg.feature_importances_
        assert fi["tree"].shape == (X.shape[1],)
        assert fi["geometry"].shape == (X.shape[1],)

    def test_geometry_importance_sums_to_one(self, fitted_regressor):
        reg, _, _ = fitted_regressor
        geom = reg.feature_importances_["geometry"]
        np.testing.assert_allclose(geom.sum(), 1.0, atol=1e-10)

    def test_get_params_set_params(self):
        reg = DirectRSRegressor(max_depth=6, alpha=0.5)
        params = reg.get_params()
        assert params["alpha"] == 0.5
        reg.set_params(alpha=10.0)
        assert reg.alpha == 10.0
