import operator
from numbers import Real

import numpy as np
import pandas as pd
import pytest

import xorq.api as xo
import xorq.expr.datatypes as dt
from xorq.common.utils.graph_utils import walk_nodes
from xorq.expr.ml.enums import ResponseMethod
from xorq.expr.ml.pipeline_lib import Pipeline, Step, make_estimator_typ
from xorq.expr.relations import Tag
from xorq.vendor.ibis.common.collections import FrozenOrderedDict
from xorq.vendor.ibis.expr.types import Expr


sklearn = pytest.importorskip("sklearn")

# sklearn submodule imports
KMeans = sklearn.cluster.KMeans
MiniBatchKMeans = sklearn.cluster.MiniBatchKMeans
RandomForestClassifier = sklearn.ensemble.RandomForestClassifier
SimpleImputer = sklearn.impute.SimpleImputer
LinearRegression = sklearn.linear_model.LinearRegression
LogisticRegression = sklearn.linear_model.LogisticRegression
accuracy_score = sklearn.metrics.accuracy_score
adjusted_rand_score = sklearn.metrics.adjusted_rand_score
f1_score = sklearn.metrics.f1_score
get_scorer = sklearn.metrics.get_scorer
get_scorer_names = sklearn.metrics.get_scorer_names
make_scorer = sklearn.metrics.make_scorer
r2_score = sklearn.metrics.r2_score
SklearnPipeline = sklearn.pipeline.Pipeline
ColumnTransformer = sklearn.compose.ColumnTransformer
SelectKBest = sklearn.feature_selection.SelectKBest
f_classif = sklearn.feature_selection.f_classif
FeatureUnion = sklearn.pipeline.FeatureUnion
OneHotEncoder = sklearn.preprocessing.OneHotEncoder
StandardScaler = sklearn.preprocessing.StandardScaler
BaseEstimator = sklearn.base.BaseEstimator
cluster = sklearn.cluster


TARGET = "target"
features = (feature0, feature1) = ("feature_0", "feature_1")


get_metadata = operator.attrgetter("metadata")


@pytest.fixture(scope="module")
def t():
    return xo.memtable(
        {
            feature0: [1, 2],
            feature1: [3, 4],
            TARGET: [0, 1],
        }
    )


step_typs = (
    sklearn.preprocessing.StandardScaler,
    sklearn.linear_model.LinearRegression,
)


@pytest.fixture(scope="module")
def sklearn_pipeline():
    sklearn_pipeline = sklearn.pipeline.make_pipeline(*(typ() for typ in step_typs))
    return sklearn_pipeline


@pytest.fixture(scope="module")
def fitted_xorq_pipeline(sklearn_pipeline, t):
    xorq_pipeline = xo.Pipeline.from_instance(sklearn_pipeline)
    return xorq_pipeline.fit(t, target=TARGET)


def test_infer_features(fitted_xorq_pipeline):
    assert all(
        features == step.features for step in fitted_xorq_pipeline.transform_steps
    )


@pytest.fixture(scope="module")
def all_tags(t, fitted_xorq_pipeline):
    expr = fitted_xorq_pipeline.predict(t)
    all_tags = walk_nodes((Tag,), expr)
    return all_tags


def test_all_tags(t, fitted_xorq_pipeline, all_tags):
    expr = fitted_xorq_pipeline.predict(t)
    actual = tuple(map(get_metadata, expr.ls.get_tags()))
    expected = tuple(map(get_metadata, all_tags))
    assert actual == expected


@pytest.mark.parametrize(
    "pairs",
    (
        (("tag", "FittedStep-transform"),),
        (("tag", "FittedStep-predict"),),
        (
            ("tag", "FittedStep-transform"),
            ("tag", "FittedStep-predict"),
        ),
    ),
)
def test_tagging_pipeline(pairs, t, fitted_xorq_pipeline):
    def contains_any_pairs(d, pairs=pairs):
        return set(pairs).intersection(d.items())

    def sort_and_tuplify(dcts):
        return tuple(sorted(tuple(sorted(dct.items())) for dct in dcts))

    actual = sort_and_tuplify(
        map(
            get_metadata,
            fitted_xorq_pipeline.predict(t).ls.get_tags(
                predicate=contains_any_pairs,
            ),
        )
    )
    expected = sort_and_tuplify(
        dct
        for dct in (
            fitted_step.tag_kwargs for fitted_step in fitted_xorq_pipeline.fitted_steps
        )
        if contains_any_pairs(dct)
    )
    assert actual and actual == expected


def test_score_expr_returns_metric(t, fitted_xorq_pipeline):
    score_expr = fitted_xorq_pipeline.score_expr(t)
    result = score_expr.execute()
    assert isinstance(result, Real)


class TestFittedStepTransform:
    """Tests for FittedStep.transform simplified logic."""

    def test_fitted_step_transform_known_schema_unpacks(self):
        """Test FittedStep.transform unpacks struct columns for known schema."""

        t = xo.memtable({"a": [1.0, 2.0, 3.0], "b": [4.0, 5.0, 6.0]})
        step = xo.Step.from_instance_name(StandardScaler(), name="scaler")
        fitted = step.fit(t, features=("a", "b"))

        result = fitted.transform(t)
        df = result.execute()

        # Should have unpacked columns a and b, not a struct column
        assert "a" in df.columns
        assert "b" in df.columns
        assert "transformed" not in df.columns

    def test_fitted_step_transform_kv_encoded_no_unpack(self):
        """Test FittedStep.transform keeps KV-encoded column without unpacking."""

        t = xo.memtable({"cat": ["x", "y", "x", "z"]})
        step = xo.Step.from_instance_name(OneHotEncoder(), name="ohe")
        fitted = step.fit(t, features=("cat",))

        result = fitted.transform(t)
        df = result.execute()

        # Should have KV-encoded column named "transformed"
        assert "transformed" in df.columns
        # Should not have unpacked category columns
        assert "cat_x" not in df.columns
        assert "cat_y" not in df.columns
        assert "cat_z" not in df.columns

    def test_fitted_step_transform_retain_others_true(self):
        """Test FittedStep.transform retains other columns by default."""

        t = xo.memtable({"a": [1.0, 2.0], "b": [3.0, 4.0], "other": ["x", "y"]})
        step = xo.Step.from_instance_name(StandardScaler(), name="scaler")
        fitted = step.fit(t, features=("a", "b"))

        result = fitted.transform(t, retain_others=True)
        df = result.execute()

        # Should retain the "other" column
        assert "other" in df.columns
        assert df["other"].tolist() == ["x", "y"]

    def test_fitted_step_transform_retain_others_false(self):
        """Test FittedStep.transform drops other columns when retain_others=False."""

        t = xo.memtable({"a": [1.0, 2.0], "b": [3.0, 4.0], "other": ["x", "y"]})
        step = xo.Step.from_instance_name(StandardScaler(), name="scaler")
        fitted = step.fit(t, features=("a", "b"))

        result = fitted.transform(t, retain_others=False)
        df = result.execute()

        # Should not retain the "other" column
        assert "other" not in df.columns


class TestPipelineGetOutputColumns:
    """Tests for Pipeline using Structer.get_output_columns."""

    def test_pipeline_known_schema_features_propagate(self):
        """Test Pipeline correctly propagates features for known schema transformers."""

        t = xo.memtable(
            {"a": [1.0, 2.0, 3.0], "b": [4.0, 5.0, 6.0], "y": [0.0, 1.0, 0.0]}
        )
        pipeline = xo.Pipeline.from_instance(
            sklearn.pipeline.make_pipeline(StandardScaler(), LinearRegression())
        )
        fitted = pipeline.fit(t, target="y")

        # The transform step should have features = ("a", "b") from the known schema
        transform_step = fitted.transform_steps[0]
        assert transform_step.structer.get_output_columns() == ("a", "b")

    def test_pipeline_kv_encoded_features_use_dest_col(self):
        """Test Pipeline correctly uses dest_col for KV-encoded transformers."""

        t = xo.memtable({"cat": ["a", "b", "a"], "y": [0.0, 1.0, 0.0]})
        pipeline = xo.Pipeline.from_instance(
            sklearn.pipeline.make_pipeline(OneHotEncoder(), LinearRegression())
        )
        fitted = pipeline.fit(t, target="y")

        # The transform step should have features = ("transformed",) for KV-encoded
        transform_step = fitted.transform_steps[0]
        assert transform_step.structer.get_output_columns("transformed") == (
            "transformed",
        )

    def test_pipeline_mixed_transform_steps(self):
        """Test Pipeline with multiple transform steps propagates features correctly."""

        t = xo.memtable(
            {"a": [1.0, None, 3.0], "b": [4.0, 5.0, 6.0], "y": [0.0, 1.0, 0.0]}
        )
        pipeline = xo.Pipeline.from_instance(
            sklearn.pipeline.make_pipeline(
                SimpleImputer(), StandardScaler(), LinearRegression()
            )
        )
        fitted = pipeline.fit(t, target="y")

        # Both transform steps should have known schema
        for transform_step in fitted.transform_steps:
            assert not transform_step.structer.is_kv_encoded

        # Prediction should work
        result = fitted.predict(t)
        assert result.execute() is not None


class TestDeeplyNestedPipelines:
    """Tests for deeply nested sklearn pipelines with xorq.

    These tests verify that xorq produces identical predictions to sklearn
    for complex nested pipeline structures.
    """

    def test_kv_encoded_deeply_nested_pipeline(self):
        """Test depth-4 nested pipeline with KV-encoded ColumnTransformer.

        Pipeline structure:
        - ColumnTransformer (KV-encoded due to OneHotEncoder)
          - FeatureUnion
            - Pipeline (SimpleImputer -> StandardScaler)
            - Pipeline (SimpleImputer -> StandardScaler)
          - Pipeline (SimpleImputer -> OneHotEncoder)
        - SelectKBest
        - RandomForestClassifier
        """

        # Create sample data
        np.random.seed(42)
        n_samples = 100

        data = pd.DataFrame(
            {
                "age": np.random.randint(18, 80, n_samples).astype(float),
                "income": np.random.randint(20000, 150000, n_samples).astype(float),
                "credit_score": np.random.randint(300, 850, n_samples).astype(float),
                "years_employed": np.random.randint(0, 40, n_samples).astype(float),
                "education": np.random.choice(
                    ["high_school", "bachelor", "master", "phd"], n_samples
                ),
                "employment_type": np.random.choice(
                    ["full_time", "part_time", "contract", "self_employed"], n_samples
                ),
                "region": np.random.choice(
                    ["north", "south", "east", "west"], n_samples
                ),
                "approved": np.random.randint(0, 2, n_samples),
            }
        )

        numeric_features = ["age", "income", "credit_score", "years_employed"]
        categorical_features = ["education", "employment_type", "region"]
        all_features = tuple(numeric_features + categorical_features)

        # Build nested sklearn pipeline
        scaled_pipeline = SklearnPipeline(
            [
                ("imputer", SimpleImputer(strategy="median")),
                ("scaler", StandardScaler()),
            ]
        )

        imputed_pipeline = SklearnPipeline(
            [
                ("imputer", SimpleImputer(strategy="mean")),
                ("scaler", StandardScaler()),
            ]
        )

        numeric_union = FeatureUnion(
            [
                ("scaled", scaled_pipeline),
                ("imputed", imputed_pipeline),
            ]
        )

        categorical_pipeline = SklearnPipeline(
            [
                ("imputer", SimpleImputer(strategy="constant", fill_value="missing")),
                (
                    "encoder",
                    OneHotEncoder(handle_unknown="ignore", sparse_output=False),
                ),
            ]
        )

        preprocessor = ColumnTransformer(
            [
                ("numeric", numeric_union, numeric_features),
                ("categorical", categorical_pipeline, categorical_features),
            ]
        )

        sklearn_pipe = SklearnPipeline(
            [
                ("preprocessor", preprocessor),
                ("selector", SelectKBest(f_classif, k=10)),
                (
                    "classifier",
                    RandomForestClassifier(n_estimators=50, random_state=42),
                ),
            ]
        )

        # Fit and predict with xorq
        expr = xo.memtable(data)
        xorq_pipeline = Pipeline.from_instance(sklearn_pipe)
        fitted_pipeline = xorq_pipeline.fit(
            expr, features=all_features, target="approved"
        )
        predictions = fitted_pipeline.predict(expr).execute()

        # Fit and predict with sklearn
        X = data[list(all_features)]
        y = data["approved"]
        sklearn_pipe.fit(X, y)
        sklearn_preds = sklearn_pipe.predict(X)

        # Assert predictions match
        assert np.array_equal(predictions[ResponseMethod.PREDICT].values, sklearn_preds)

    def test_non_kv_deeply_nested_pipeline(self):
        """Test depth-4 nested pipeline with all known-schema transformers.

        Pipeline structure:
        - ColumnTransformer (known schema - no KV-encoded children)
          - Pipeline (SimpleImputer -> StandardScaler -> Pipeline)
            - Pipeline (SimpleImputer -> StandardScaler)
          - Pipeline (SimpleImputer -> StandardScaler)
        - RandomForestClassifier
        """

        # Create sample data
        np.random.seed(42)
        n_samples = 100

        data = pd.DataFrame(
            {
                "age": np.random.randint(18, 80, n_samples).astype(float),
                "income": np.random.randint(20000, 150000, n_samples).astype(float),
                "credit_score": np.random.randint(300, 850, n_samples).astype(float),
                "years_employed": np.random.randint(0, 40, n_samples).astype(float),
                "debt_ratio": np.random.uniform(0, 1, n_samples),
                "savings": np.random.randint(0, 100000, n_samples).astype(float),
                "approved": np.random.randint(0, 2, n_samples),
            }
        )

        numeric_features_a = ["age", "income", "credit_score"]
        numeric_features_b = ["years_employed", "debt_ratio", "savings"]
        all_features = tuple(numeric_features_a + numeric_features_b)

        # Build nested sklearn pipeline (depth 4)
        inner_pipeline = SklearnPipeline(
            [
                ("imputer2", SimpleImputer(strategy="mean")),
                ("scaler2", StandardScaler()),
            ]
        )

        numeric_a_pipeline = SklearnPipeline(
            [
                ("imputer", SimpleImputer(strategy="median")),
                ("scaler", StandardScaler()),
                ("inner", inner_pipeline),
            ]
        )

        numeric_b_pipeline = SklearnPipeline(
            [
                ("imputer", SimpleImputer(strategy="median")),
                ("scaler", StandardScaler()),
            ]
        )

        preprocessor = ColumnTransformer(
            [
                ("numeric_a", numeric_a_pipeline, numeric_features_a),
                ("numeric_b", numeric_b_pipeline, numeric_features_b),
            ]
        )

        sklearn_pipe = SklearnPipeline(
            [
                ("preprocessor", preprocessor),
                (
                    "classifier",
                    RandomForestClassifier(n_estimators=50, random_state=42),
                ),
            ]
        )

        # Fit and predict with xorq
        expr = xo.memtable(data)
        xorq_pipeline = Pipeline.from_instance(sklearn_pipe)
        fitted_pipeline = xorq_pipeline.fit(
            expr, features=all_features, target="approved"
        )
        predictions = fitted_pipeline.predict(expr).execute()

        # Fit and predict with sklearn
        X = data[list(all_features)]
        y = data["approved"]
        sklearn_pipe.fit(X, y)
        sklearn_preds = sklearn_pipe.predict(X)

        # Assert predictions match
        assert np.array_equal(predictions[ResponseMethod.PREDICT].values, sklearn_preds)


def _scorer_info():
    """Return [(name, module, response_method), ...] for every registered scorer."""

    return [
        (
            name,
            get_scorer(name)._score_func.__module__,
            get_scorer(name)._response_method,
        )
        for name in get_scorer_names()
    ]


def get_scorers_by_type():
    """Categorize all sklearn scorers by their type based on internal module path."""
    info = _scorer_info()

    proba = {
        name
        for name, _, response in info
        if isinstance(response, tuple)
        or response in (ResponseMethod.PREDICT_PROBA, ResponseMethod.DECISION_FUNCTION)
    }

    multilabel = {name for name, _, _ in info if name.endswith("_samples")}

    # Non-multilabel scorers categorized by module
    non_ml = [
        (name, module) for name, module, _ in info if not name.endswith("_samples")
    ]

    classification = {
        name
        for name, module in non_ml
        if "_classification" in module or "_ranking" in module or "_scorer" in module
    }

    regression = {name for name, module in non_ml if "_regression" in module}

    cluster = {name for name, module in non_ml if "cluster" in module}

    return FrozenOrderedDict(
        **{
            k: tuple(sorted(v))
            for k, v in (
                ("classification", classification),
                ("regression", regression),
                ("cluster", cluster),
                ("proba", proba),
                ("multilabel", multilabel),
            )
        }
    )


SCORERS_BY_TYPE = get_scorers_by_type()


class TestPipelineScoringMatchSklearn:
    """Tests for pipeline scoring with all compatible scorers."""

    @pytest.fixture
    def scoring_data(self):
        """Generate dataset suitable for classification, regression, and clustering."""

        np.random.seed(42)
        n = 100
        return {
            "x1": np.random.randn(n).tolist(),
            "x2": np.random.randn(n).tolist(),
            "y_class": (np.random.randn(n) > 0).astype(int).tolist(),
            "y_reg": (
                np.abs(np.random.randn(n)) + 0.1
            ).tolist(),  # positive for log/deviance scorers
        }

    @pytest.mark.parametrize("scorer_name", SCORERS_BY_TYPE["classification"])
    def test_classifier_scorer(self, scoring_data, scorer_name):
        """Test classification scorers match sklearn."""

        X = np.array([scoring_data["x1"], scoring_data["x2"]]).T
        y = np.array(scoring_data["y_class"])

        sklearn_pipe = SklearnPipeline(
            [
                ("scaler", StandardScaler()),
                ("model", LogisticRegression(random_state=42, max_iter=1000)),
            ]
        )
        sklearn_pipe.fit(X, y)

        t = xo.memtable(scoring_data)
        xorq_pipeline = xo.Pipeline.from_instance(sklearn_pipe)
        fitted_xorq = xorq_pipeline.fit(t, features=("x1", "x2"), target="y_class")

        scorer = get_scorer(scorer_name)
        sklearn_score = scorer(sklearn_pipe, X, y)
        xorq_score = fitted_xorq.score(X, y, scorer=scorer_name)

        np.testing.assert_allclose(xorq_score, sklearn_score, rtol=1e-9, atol=1e-12)

    @pytest.mark.parametrize("scorer_name", SCORERS_BY_TYPE["regression"])
    def test_regressor_scorer(self, scoring_data, scorer_name):
        """Test regression scorers match sklearn."""

        X = np.array([scoring_data["x1"], scoring_data["x2"]]).T
        y = np.array(scoring_data["y_reg"])

        sklearn_pipe = SklearnPipeline(
            [("scaler", StandardScaler()), ("model", LinearRegression())]
        )
        sklearn_pipe.fit(X, y)

        t = xo.memtable(scoring_data)
        xorq_pipeline = xo.Pipeline.from_instance(sklearn_pipe)
        fitted_xorq = xorq_pipeline.fit(t, features=("x1", "x2"), target="y_reg")

        scorer = get_scorer(scorer_name)
        sklearn_score = scorer(sklearn_pipe, X, y)
        xorq_score = fitted_xorq.score(X, y, scorer=scorer_name)

        np.testing.assert_allclose(xorq_score, sklearn_score, rtol=1e-9, atol=1e-12)

    @pytest.mark.parametrize("scorer_name", SCORERS_BY_TYPE["cluster"])
    def test_cluster_scorer(self, scoring_data, scorer_name):
        """Test clustering scorers match sklearn."""

        X = np.array([scoring_data["x1"], scoring_data["x2"]]).T
        y = np.array(scoring_data["y_class"])

        sklearn_pipe = SklearnPipeline(
            [
                ("scaler", StandardScaler()),
                ("clusterer", KMeans(n_clusters=2, random_state=42, n_init=10)),
            ]
        )
        sklearn_pipe.fit(X, y)

        t = xo.memtable(scoring_data)
        xorq_pipeline = xo.Pipeline.from_instance(sklearn_pipe)
        fitted_xorq = xorq_pipeline.fit(t, features=("x1", "x2"), target="y_class")

        scorer = get_scorer(scorer_name)
        sklearn_score = scorer(sklearn_pipe, X, y)
        xorq_score = fitted_xorq.score(X, y, scorer=scorer_name)

        np.testing.assert_allclose(xorq_score, sklearn_score, rtol=1e-9, atol=1e-12)


class TestScoreExpr:
    """Tests for FittedPipeline.score_expr and .score edge cases."""

    @pytest.fixture
    def fitted_classifier(self):
        """Fitted classifier pipeline."""

        np.random.seed(42)
        data = {"x1": [0.0, 1.0], "x2": [0.0, 1.0], "y": [0, 1]}
        X = np.array([data["x1"], data["x2"]]).T
        y = np.array(data["y"])

        sklearn_pipe = SklearnPipeline(
            [("scaler", StandardScaler()), ("model", LogisticRegression())]
        )
        sklearn_pipe.fit(X, y)

        t = xo.memtable(data)
        fitted = xo.Pipeline.from_instance(sklearn_pipe).fit(
            t, features=("x1", "x2"), target="y"
        )
        return fitted, sklearn_pipe, X, y, t

    @pytest.fixture
    def fitted_regressor(self):
        """Fitted regressor pipeline."""

        np.random.seed(42)
        data = {"x1": [0.0, 1.0], "x2": [0.0, 1.0], "y": [0.0, 1.0]}
        X = np.array([data["x1"], data["x2"]]).T
        y = np.array(data["y"])

        sklearn_pipe = SklearnPipeline(
            [("scaler", StandardScaler()), ("model", LinearRegression())]
        )
        sklearn_pipe.fit(X, y)

        t = xo.memtable(data)
        fitted = xo.Pipeline.from_instance(sklearn_pipe).fit(
            t, features=("x1", "x2"), target="y"
        )
        return fitted, sklearn_pipe, X, y, t

    def test_default_scorer_classifier_is_accuracy(self, fitted_classifier):
        """Test default scorer for classifier is accuracy_score."""

        fitted, *_ = fitted_classifier
        scorer = fitted._get_default_scorer()
        assert scorer._score_func is accuracy_score

    def test_default_scorer_regressor_is_r2(self, fitted_regressor):
        """Test default scorer for regressor is r2_score."""

        fitted, *_ = fitted_regressor
        scorer = fitted._get_default_scorer()
        assert scorer._score_func is r2_score

    def test_default_scorer_cluster_is_adjusted_rand(self):
        """Test default scorer for clustering is adjusted_rand_score."""

        t = xo.memtable({"x": [0.0, 1.0], "y": [0, 1]})
        fitted = xo.Pipeline.from_instance(
            sklearn.pipeline.make_pipeline(KMeans(n_clusters=2, n_init=1))
        ).fit(t, features=("x",), target="y")

        assert fitted._get_default_scorer()._score_func is adjusted_rand_score

    def test_string_scorer(self, fitted_classifier):
        """Test passing a scorer name string."""

        fitted, sklearn_pipe, X, y, _ = fitted_classifier

        xorq_score = fitted.score(X, y, scorer="f1")
        sklearn_score = get_scorer("f1")(sklearn_pipe, X, y)
        np.testing.assert_allclose(xorq_score, sklearn_score, rtol=1e-9)

    def test_callable_scorer(self, fitted_classifier):
        """Test passing a raw callable metric function."""

        fitted, sklearn_pipe, X, y, _ = fitted_classifier

        xorq_score = fitted.score(X, y, scorer=f1_score)
        sklearn_score = f1_score(y, sklearn_pipe.predict(X))
        np.testing.assert_allclose(xorq_score, sklearn_score, rtol=1e-9)

    def test_make_scorer_object(self, fitted_classifier):
        """Test passing a make_scorer object directly."""

        fitted, sklearn_pipe, X, y, _ = fitted_classifier

        scorer = make_scorer(f1_score)
        xorq_score = fitted.score(X, y, scorer=scorer)
        sklearn_score = scorer(sklearn_pipe, X, y)
        np.testing.assert_allclose(xorq_score, sklearn_score, rtol=1e-9)

    def test_score_expr_returns_expression(self, fitted_classifier):
        """Test score_expr returns an ibis expression."""

        fitted, _, _, _, t = fitted_classifier
        expr = fitted.score_expr(t, scorer="accuracy")

        assert isinstance(expr, Expr)
        assert isinstance(expr.execute(), (int, float))

    def test_default_scorer_raises_for_unknown_model(self):
        """Test _get_default_scorer raises ValueError for unknown model type."""

        # Custom estimator that isn't a Classifier/Regressor/Cluster
        class CustomEstimator(BaseEstimator):
            return_type = dt.int64  # needed for xorq to handle predict

            def fit(self, X, y=None):
                return self

            def predict(self, X):
                return [0] * len(X)

        t = xo.memtable({"x": [0.0, 1.0], "y": [0, 1]})
        fitted = xo.Pipeline.from_instance(
            sklearn.pipeline.make_pipeline(CustomEstimator())
        ).fit(t, features=("x",), target="y")

        with pytest.raises(ValueError, match="Cannot determine default scorer"):
            fitted._get_default_scorer()


class TestStepFromFitFunctions:
    """Tests for Step.from_fit_transform and Step.from_fit_predict.

    These exercise make_estimator_typ which dynamically creates BaseEstimator
    subclasses from raw fit/transform or fit/predict callables.
    """

    def test_from_fit_transform_creates_transform_type(self):
        """Step.from_fit_transform creates a type with transform (not predict)."""

        def my_fit(X, y=None):
            return np.mean(X, axis=0)

        def my_transform(model, X, y=None):
            return X - model

        step = Step.from_fit_transform(
            fit=my_fit,
            transform=my_transform,
            return_type=dt.Array(dt.float64),
            name="custom_transform",
        )

        # Verify the dynamically created type has correct attributes
        assert hasattr(step.instance, "transform")
        assert hasattr(step.instance, "fit")
        assert not hasattr(step.instance, "predict")
        assert step.instance.return_type == dt.Array(dt.float64)

    def test_from_fit_predict_creates_predict_type(self):
        """Step.from_fit_predict creates a type with predict (not transform)."""

        def my_fit(X, y=None):
            return int(np.median(y))

        def my_predict(model, X, y=None):
            return np.full(len(X), model)

        step = Step.from_fit_predict(
            fit=my_fit,
            predict=my_predict,
            return_type=dt.int64,
            name="custom_predict",
        )

        assert hasattr(step.instance, "predict")
        assert hasattr(step.instance, "fit")
        assert not hasattr(step.instance, "transform")
        assert step.instance.return_type == dt.int64

    def test_make_estimator_typ_both_raises(self):
        """Passing both transform and predict raises ValueError."""

        with pytest.raises(ValueError):
            make_estimator_typ(
                fit=lambda X, y=None: None,
                return_type=dt.float64,
                transform=lambda m, X: X,
                predict=lambda m, X: X,
            )

    def test_make_estimator_typ_neither_raises(self):
        """Passing neither transform nor predict raises ValueError."""

        with pytest.raises(ValueError):
            make_estimator_typ(
                fit=lambda X, y=None: None,
                return_type=dt.float64,
            )

    def test_from_fit_predict_end_to_end(self):
        """Step.from_fit_predict works end-to-end with dest_col."""

        def my_fit(X, y=None):
            return int(np.median(y))

        def my_predict(model, X, y=None):
            return np.full(len(X), model)

        step = Step.from_fit_predict(
            fit=my_fit,
            predict=my_predict,
            return_type=dt.int64,
            name="custom_predict",
        )

        t = xo.memtable({"a": [1.0, 2.0, 3.0], "b": [4.0, 5.0, 6.0], "y": [0, 1, 1]})
        fitted = step.fit(t, features=("a", "b"), target="y", dest_col="pred")
        result = fitted.predict(t)
        df = result.execute()
        assert df is not None
        assert len(df) == 3


class TestFeatureImportances:
    """Tests for FittedStep.feature_importances and FittedPipeline.feature_importances."""

    def test_fitted_step_feature_importances(self):
        """FittedStep.feature_importances returns importances for tree models."""

        t = xo.memtable(
            {
                "a": np.random.randn(50).tolist(),
                "b": np.random.randn(50).tolist(),
                "y": (np.random.randn(50) > 0).astype(int).tolist(),
            }
        )

        step = xo.Step.from_instance_name(
            RandomForestClassifier(n_estimators=5, random_state=42),
            name="rf",
        )
        fitted = step.fit(t, features=("a", "b"), target="y")
        result = fitted.feature_importances(t)
        df = result.execute()

        assert df is not None
        assert "feature_importances" in df.columns
        importances = df["feature_importances"].iloc[0]
        assert len(importances) == 2  # two features
        assert all(isinstance(v, float) for v in importances)

    def test_fitted_pipeline_feature_importances(self):
        """FittedPipeline.feature_importances returns importances through pipeline."""

        t = xo.memtable(
            {
                "a": np.random.randn(50).tolist(),
                "b": np.random.randn(50).tolist(),
                "y": (np.random.randn(50) > 0).astype(int).tolist(),
            }
        )

        sklearn_pipe = SklearnPipeline(
            [
                ("scaler", StandardScaler()),
                ("rf", RandomForestClassifier(n_estimators=5, random_state=42)),
            ]
        )
        fitted = Pipeline.from_instance(sklearn_pipe).fit(
            t, features=("a", "b"), target="y"
        )

        result = fitted.feature_importances(t)
        df = result.execute()

        assert df is not None
        assert "feature_importances" in df.columns
        importances = df["feature_importances"].iloc[0]
        assert len(importances) == 2


class TestClusteringPredict:
    """Tests for clustering algorithm predict support."""

    @pytest.fixture
    def cluster_data(self):
        """Generate data with clear cluster structure."""

        np.random.seed(42)
        # Two well-separated clusters
        cluster1 = np.random.randn(10, 2) + [0, 0]
        cluster2 = np.random.randn(10, 2) + [10, 10]
        data = np.vstack([cluster1, cluster2])
        return {"num1": data[:, 0].tolist(), "num2": data[:, 1].tolist()}

    @pytest.mark.parametrize(
        "clusterer_cls,clusterer_kwargs",
        [
            pytest.param(
                "KMeans",
                {"n_clusters": 2, "random_state": 42, "n_init": 10},
                id="KMeans",
            ),
            pytest.param(
                "MiniBatchKMeans",
                {"n_clusters": 2, "random_state": 42, "n_init": 10},
                id="MiniBatchKMeans",
            ),
            pytest.param(
                "BisectingKMeans",
                {"n_clusters": 2, "random_state": 42},
                id="BisectingKMeans",
            ),
            pytest.param(
                "Birch",
                {"n_clusters": 2},
                id="Birch",
            ),
            pytest.param(
                "MeanShift",
                {},
                id="MeanShift",
            ),
            pytest.param(
                "AffinityPropagation",
                {"random_state": 42},
                id="AffinityPropagation",
            ),
        ],
    )
    def test_inductive_clustering_predict(
        self, cluster_data, clusterer_cls, clusterer_kwargs
    ):
        """Test that inductive clustering algorithms support predict."""

        t = xo.memtable(cluster_data)
        features = ("num1", "num2")

        ClustererClass = getattr(cluster, clusterer_cls)
        clusterer = ClustererClass(**clusterer_kwargs)

        # xorq predict
        step = xo.Step.from_instance_name(clusterer, name="clusterer")
        fitted = step.fit(t, features=features)
        result = fitted.predict(t)
        xorq_labels = result.execute()[ResponseMethod.PREDICT].values

        # sklearn predict
        X = np.array([cluster_data["num1"], cluster_data["num2"]]).T
        sklearn_clusterer = ClustererClass(**clusterer_kwargs)
        sklearn_clusterer.fit(X)
        sklearn_labels = sklearn_clusterer.predict(X)

        # Labels should match
        np.testing.assert_array_equal(xorq_labels, sklearn_labels)

    @pytest.mark.parametrize(
        "clusterer_cls,clusterer_kwargs",
        [
            pytest.param(
                "DBSCAN",
                {"eps": 3, "min_samples": 2},
                id="DBSCAN",
            ),
            pytest.param(
                "HDBSCAN",
                {"min_samples": 2},
                id="HDBSCAN",
            ),
            pytest.param(
                "AgglomerativeClustering",
                {"n_clusters": 2},
                id="AgglomerativeClustering",
            ),
            pytest.param(
                "SpectralClustering",
                {"n_clusters": 2, "random_state": 42},
                id="SpectralClustering",
            ),
            pytest.param(
                "OPTICS",
                {"min_samples": 2},
                id="OPTICS",
            ),
        ],
    )
    def test_transductive_clustering_rejected_at_fit(
        self, cluster_data, clusterer_cls, clusterer_kwargs
    ):
        """Test that transductive clustering algorithms are rejected at fit time."""

        t = xo.memtable(cluster_data)
        features = ("num1", "num2")

        ClustererClass = getattr(cluster, clusterer_cls)
        clusterer = ClustererClass(**clusterer_kwargs)

        step = xo.Step.from_instance_name(clusterer, name="clusterer")

        with pytest.raises(ValueError, match="must have transform or predict method"):
            step.fit(t, features=features)

    def test_pipeline_fit_without_target_for_clustering(self, cluster_data):
        """Test Pipeline.fit allows ClusterMixin predict steps without a target."""

        t = xo.memtable(cluster_data)
        features = ("num1", "num2")

        sklearn_pipe = SklearnPipeline(
            [
                ("scaler", StandardScaler()),
                (
                    "clusterer",
                    MiniBatchKMeans(n_clusters=2, random_state=42, n_init=10),
                ),
            ]
        )

        xorq_pipeline = Pipeline.from_instance(sklearn_pipe)
        fitted = xorq_pipeline.fit(t, features=features)

        # Should be able to predict without ever providing a target
        result = fitted.predict(t)
        xorq_labels = result.execute()[ResponseMethod.PREDICT].values

        # Verify against sklearn
        X = np.array([cluster_data["num1"], cluster_data["num2"]]).T
        sklearn_pipe.fit(X)
        sklearn_labels = sklearn_pipe.predict(X)

        np.testing.assert_array_equal(xorq_labels, sklearn_labels)
