"""Tests for scikit-learn API compatibility."""

import numpy as np
import pytest
from sklearn.utils.estimator_checks import parametrize_with_checks
from sklearn.pipeline import Pipeline
from sklearn.model_selection import cross_val_score
from sklearn.preprocessing import StandardScaler

from surrogate_forest import (
    SurrogateDecisionTreeClassifier,
    SurrogateDecisionTreeRegressor,
    SurrogateRandomForestClassifier,
    SurrogateRandomForestRegressor,
)


@parametrize_with_checks([
    SurrogateDecisionTreeClassifier(random_state=42),
    SurrogateDecisionTreeRegressor(random_state=42),
    SurrogateRandomForestClassifier(n_estimators=5, random_state=42),
    SurrogateRandomForestRegressor(n_estimators=5, random_state=42),
])
def test_sklearn_compatible(estimator, check):
    # Bootstrap forests can't achieve exact weight equivalence because
    # bootstrap sampling from n weighted samples vs n*k duplicated samples
    # produces fundamentally different bootstrap datasets
    check_name = check.func.__name__ if hasattr(check, "func") else str(check)
    is_forest = "RandomForest" in type(estimator).__name__
    is_weight_equiv = "sample_weight_equivalence" in str(check_name)
    if is_forest and is_weight_equiv:
        pytest.xfail("Bootstrap forests: weight/duplication equivalence not exact")
    check(estimator)


class TestPipelineIntegration:
    def test_classifier_in_pipeline(self, clean_classification_data):
        X, y = clean_classification_data
        pipe = Pipeline([
            ("scaler", StandardScaler()),
            ("clf", SurrogateRandomForestClassifier(
                n_estimators=5, max_depth=3, random_state=42
            )),
        ])
        pipe.fit(X, y)
        preds = pipe.predict(X)
        assert len(preds) == len(y)

    def test_cross_val_score(self, clean_classification_data):
        X, y = clean_classification_data
        clf = SurrogateRandomForestClassifier(
            n_estimators=5, max_depth=3, random_state=42
        )
        scores = cross_val_score(clf, X, y, cv=3)
        assert len(scores) == 3
        assert all(0 <= s <= 1 for s in scores)


class TestCloneAndParams:
    def test_get_params(self):
        clf = SurrogateDecisionTreeClassifier(max_depth=5, max_surrogates=3)
        params = clf.get_params()
        assert params["max_depth"] == 5
        assert params["max_surrogates"] == 3

    def test_set_params(self):
        clf = SurrogateDecisionTreeClassifier()
        clf.set_params(max_depth=10)
        assert clf.max_depth == 10

    def test_clone(self):
        from sklearn.base import clone
        clf = SurrogateRandomForestClassifier(n_estimators=50, max_surrogates=3)
        clf2 = clone(clf)
        assert clf2.n_estimators == 50
        assert clf2.max_surrogates == 3
        assert not hasattr(clf2, "estimators_")
