"""Tests for decision tree estimators."""

import numpy as np
import pytest

from surrogate_forest import (
    SurrogateDecisionTreeClassifier,
    SurrogateDecisionTreeRegressor,
)


class TestClassifierTree:
    def test_fit_predict_clean(self, clean_classification_data):
        X, y = clean_classification_data
        clf = SurrogateDecisionTreeClassifier(max_depth=5, random_state=42)
        clf.fit(X, y)
        acc = (clf.predict(X) == y).mean()
        assert acc > 0.85

    def test_fit_predict_missing(self, missing_classification_data):
        X, y = missing_classification_data
        clf = SurrogateDecisionTreeClassifier(max_depth=5, random_state=42)
        clf.fit(X, y)
        preds = clf.predict(X)
        acc = (preds == y).mean()
        assert acc > 0.7

    def test_predict_proba(self, clean_classification_data):
        X, y = clean_classification_data
        clf = SurrogateDecisionTreeClassifier(max_depth=3, random_state=42)
        clf.fit(X, y)
        proba = clf.predict_proba(X)
        assert proba.shape == (len(X), 2)
        assert np.allclose(proba.sum(axis=1), 1.0)

    def test_multiclass(self, multiclass_data):
        X, y = multiclass_data
        clf = SurrogateDecisionTreeClassifier(max_depth=5, random_state=42)
        clf.fit(X, y)
        assert clf.n_classes_ == 3
        proba = clf.predict_proba(X)
        assert proba.shape[1] == 3

    def test_attributes(self, clean_classification_data):
        X, y = clean_classification_data
        clf = SurrogateDecisionTreeClassifier(random_state=42)
        clf.fit(X, y)
        assert clf.n_features_in_ == 5
        assert hasattr(clf, "tree_")
        assert hasattr(clf, "classes_")

    def test_symmetric_mode(self, clean_classification_data):
        X, y = clean_classification_data
        clf = SurrogateDecisionTreeClassifier(
            max_depth=3, symmetric_tree=True, random_state=42
        )
        clf.fit(X, y)
        preds = clf.predict(X)
        assert len(np.unique(preds)) >= 2

    def test_best_first_mode(self, clean_classification_data):
        X, y = clean_classification_data
        clf = SurrogateDecisionTreeClassifier(
            max_leaf_nodes=8, random_state=42
        )
        clf.fit(X, y)
        acc = (clf.predict(X) == y).mean()
        assert acc > 0.7

    def test_exact_mode(self, clean_classification_data):
        X, y = clean_classification_data
        clf = SurrogateDecisionTreeClassifier(
            max_depth=4, use_histogram=False, random_state=42
        )
        clf.fit(X, y)
        acc = (clf.predict(X) == y).mean()
        assert acc > 0.8


class TestRegressorTree:
    def test_fit_predict_clean(self, clean_regression_data):
        X, y = clean_regression_data
        reg = SurrogateDecisionTreeRegressor(max_depth=5, random_state=42)
        reg.fit(X, y)
        from sklearn.metrics import r2_score
        assert r2_score(y, reg.predict(X)) > 0.5

    def test_fit_predict_missing(self, missing_regression_data):
        X, y = missing_regression_data
        reg = SurrogateDecisionTreeRegressor(max_depth=5, random_state=42)
        reg.fit(X, y)
        preds = reg.predict(X)
        assert not np.any(np.isnan(preds))

    def test_mae_criterion(self, clean_regression_data):
        X, y = clean_regression_data
        reg = SurrogateDecisionTreeRegressor(
            criterion="absolute_error", max_depth=4,
            use_histogram=False, random_state=42
        )
        reg.fit(X, y)
        preds = reg.predict(X)
        assert not np.any(np.isnan(preds))
