"""Tests for ensemble estimators."""

import numpy as np
import pytest

from surrogate_forest import (
    SurrogateRandomForestClassifier,
    SurrogateRandomForestRegressor,
)


class TestForestClassifier:
    def test_fit_predict(self, clean_classification_data):
        X, y = clean_classification_data
        rf = SurrogateRandomForestClassifier(
            n_estimators=20, max_depth=5, random_state=42
        )
        rf.fit(X, y)
        acc = (rf.predict(X) == y).mean()
        assert acc > 0.85

    def test_predict_proba(self, clean_classification_data):
        X, y = clean_classification_data
        rf = SurrogateRandomForestClassifier(
            n_estimators=10, max_depth=3, random_state=42
        )
        rf.fit(X, y)
        proba = rf.predict_proba(X)
        assert proba.shape == (len(X), 2)
        assert np.allclose(proba.sum(axis=1), 1.0, atol=1e-6)

    def test_oob_score(self, clean_classification_data):
        X, y = clean_classification_data
        rf = SurrogateRandomForestClassifier(
            n_estimators=50, max_depth=5, oob_score=True, random_state=42
        )
        rf.fit(X, y)
        assert hasattr(rf, "oob_score_")
        assert 0.0 <= rf.oob_score_ <= 1.0

    def test_missing_data(self, missing_classification_data):
        X, y = missing_classification_data
        rf = SurrogateRandomForestClassifier(
            n_estimators=20, max_depth=5, random_state=42
        )
        rf.fit(X, y)
        preds = rf.predict(X)
        acc = (preds == y).mean()
        assert acc > 0.7

    def test_n_estimators(self, clean_classification_data):
        X, y = clean_classification_data
        rf = SurrogateRandomForestClassifier(n_estimators=5, random_state=42)
        rf.fit(X, y)
        assert len(rf.estimators_) == 5

    def test_no_bootstrap(self, clean_classification_data):
        X, y = clean_classification_data
        rf = SurrogateRandomForestClassifier(
            n_estimators=5, bootstrap=False, random_state=42
        )
        rf.fit(X, y)
        assert (rf.predict(X) == y).mean() > 0.5

    def test_parallel(self, clean_classification_data):
        X, y = clean_classification_data
        rf = SurrogateRandomForestClassifier(
            n_estimators=10, n_jobs=2, random_state=42
        )
        rf.fit(X, y)
        assert len(rf.estimators_) == 10


class TestForestRegressor:
    def test_fit_predict(self, clean_regression_data):
        X, y = clean_regression_data
        rf = SurrogateRandomForestRegressor(
            n_estimators=20, max_depth=5, random_state=42
        )
        rf.fit(X, y)
        from sklearn.metrics import r2_score
        assert r2_score(y, rf.predict(X)) > 0.5

    def test_oob_score(self, clean_regression_data):
        X, y = clean_regression_data
        rf = SurrogateRandomForestRegressor(
            n_estimators=50, max_depth=5, oob_score=True, random_state=42
        )
        rf.fit(X, y)
        assert hasattr(rf, "oob_score_")

    def test_missing_data(self, missing_regression_data):
        X, y = missing_regression_data
        rf = SurrogateRandomForestRegressor(
            n_estimators=20, max_depth=5, random_state=42
        )
        rf.fit(X, y)
        preds = rf.predict(X)
        assert not np.any(np.isnan(preds))
