scikit-learn Transformers Implementation Plan¶

import narwhals as nw
import numpy as np
import polars as pl
import pytest

from tusk.exceptions import SchemaError
from tusk.sklearn._frames import as_keys, collect_matrix


@pytest.mark.parametrize(
    ("label", "value"),
    [
        ("list", [1, 2, 3]),
        ("range", range(1, 4)),
        ("1d array", np.array([1, 2, 3])),
        ("column array", np.array([[1], [2], [3]])),
        ("one column frame", pl.DataFrame({"id": [1, 2, 3]})),
    ],
)
def test_as_keys_accepts_every_documented_shape(label, value):
    assert as_keys(value) == [1, 2, 3]


@pytest.mark.parametrize(
    "value",
    [np.array([[1, 2], [3, 4]]), pl.DataFrame({"a": [1], "b": [2]})],
)
def test_as_keys_rejects_more_than_one_column(value):
    with pytest.raises(SchemaError, match="exactly one column"):
        as_keys(value)


def test_collect_matrix_returns_rows_in_key_order():
    frame = pl.LazyFrame({"id": [1, 2, 3, 4], "a": [10.0, 20.0, 30.0, 40.0]})
    out = nw.from_native(collect_matrix(frame, "id", [3, 1, 4], None))
    assert out["a"].to_list() == [30.0, 10.0, 40.0]


def test_collect_matrix_drops_the_primary_key():
    frame = pl.LazyFrame({"id": [1, 2], "a": [10.0, 20.0]})
    out = nw.from_native(collect_matrix(frame, "id", [1, 2], None))
    assert out.columns == ["a"]


def test_collect_matrix_rejects_a_duplicate_key():
    frame = pl.LazyFrame({"id": [1, 2], "a": [10.0, 20.0]})
    with pytest.raises(SchemaError, match="duplicate"):
        collect_matrix(frame, "id", [1, 1], None)


def test_collect_matrix_rejects_a_key_with_no_row():
    frame = pl.LazyFrame({"id": [1, 2], "a": [10.0, 20.0]})
    with pytest.raises(SchemaError, match="no row"):
        collect_matrix(frame, "id", [1, 99], None)


def test_backend_hint_reraises_the_original_exception_type():
    from tusk.sklearn._frames import backend_hint

    with pytest.raises(ValueError, match="boom"):
        with backend_hint(pl.DataFrame({"a": [1]})):
            raise ValueError("boom")


def test_backend_hint_names_the_backend_and_the_fix():
    import sys

    from tusk.sklearn._frames import backend_hint

    try:
        with backend_hint(pl.DataFrame({"a": [1]})):
            raise ValueError("boom")
    except ValueError as exc:
        hint = " ".join(getattr(exc, "__notes__", []))
        if sys.version_info < (3, 11):
            return
    assert "polars" in hint.lower()
    assert "output_backend" in hint

class LineageError(TuskError):
    """Raised when a kept encoded column has no counterpart after the refit.

    The only way sentinel lineage can *under*-keep is by missing a source, so
    this is the tripwire for that case: if a feature was wrongly pruned, a name
    the frozen mask needs goes missing from the refitted encoder's output.
    Failing here is the alternative to silently serving different columns than
    the ones the selector chose.
    """


class EncoderError(TuskError):
    """Raised when the supplied encoder cannot be refit on a column subset.

    Pruning narrows the matrix, so the encoder is refit on fewer columns than
    it first saw. A ``ColumnTransformer`` naming its columns explicitly cannot
    survive that -- and after DFS it is already a mistake for two other
    reasons, so it is refused rather than worked around.
    """


class UnencodedFeatureWarning(UserWarning):
    """Warns that a feature fed no encoded column, so it was pruned.

    The encoder simply never looked at it -- most often a ``ColumnTransformer``
    covering only some dtypes while ``remainder`` stays at its ``"drop"``
    default. Pruning it is self-consistent, but doing so silently would let a
    user lose every numeric feature without a word. Its own class, so it can be
    filtered independently.
    """


class LineageWarning(UserWarning):
    """Warns that lineage was unrecoverable, so nothing was pruned.

    An encoder whose output names do not mention their inputs -- ``PCA`` gives
    ``pca0`` -- severs the link between kept columns and tusk features. Keeping
    every feature is the correct fallback, since pruning is an optimization and
    no result depends on it. Its own class, so it can be filtered independently.
    """

"""scikit-learn estimators wrapping deep feature synthesis.

Deliberately not imported by :mod:`tusk`: importing scikit-learn eagerly would
break every install without the ``sklearn`` extra. Reach these through
``from tusk.sklearn import DFSTransformer``, as with ``xgboost.sklearn``.
"""

from __future__ import annotations

__all__: list[str] = []

"""The collect boundary.

tusk never collects; this is the one module that does, because scikit-learn
cannot consume a query plan. Everything here is about crossing that line
exactly once, in a way that cannot silently misalign rows against ``y``.
"""

from __future__ import annotations

import contextlib
import warnings
from collections.abc import Iterator, Sequence
from typing import Any

import narwhals as nw
import numpy as np

from tusk.exceptions import SchemaError

_POSITION = "__tusk_position"


def as_keys(X: Any) -> list[Any]:
    """Normalize ``X`` to a flat list of primary-key values.

    Args:
        X: A 1-D sequence, an ``(n, 1)`` array, or a one-column frame.

    Returns:
        The key values, in the order given. That order is the contract: it
        becomes the row order of the feature matrix, which is what keeps ``X``
        aligned with ``y``.

    Raises:
        SchemaError: If ``X`` has more than one column, or more than two
            dimensions.
    """
    frame = nw.from_native(X, eager_only=True, pass_through=True)
    if isinstance(frame, nw.DataFrame):
        if len(frame.columns) != 1:
            raise SchemaError(
                f"X must have exactly one column, the target's primary key; "
                f"got {len(frame.columns)}",
            )
        return frame[frame.columns[0]].to_list()
    array = np.asarray(X)
    if array.ndim == 2:
        if array.shape[1] != 1:
            raise SchemaError(
                f"X must have exactly one column, the target's primary key; "
                f"got {array.shape[1]}",
            )
        array = array[:, 0]
    elif array.ndim != 1:
        raise SchemaError(f"X must be 1- or 2-dimensional; got {array.ndim}")
    return array.tolist()


def collect_matrix(
    lazy_native: Any,
    primary_key: str,
    keys: Sequence[Any],
    output_backend: str | None,
) -> Any:
    """Materialize the feature matrix for ``keys``, in ``keys`` order.

    Args:
        lazy_native: The uncomputed matrix from ``apply_features``.
        primary_key: The target table's key column.
        keys: Key values selecting and ordering the rows.
        output_backend: Backend to collect to, or None to collect natively.

    Returns:
        An eager native frame with one row per key, in key order, without the
        primary key column -- it is a join key, not a feature.

    Raises:
        SchemaError: If ``keys`` repeats a value, or names a key that produced
            no row.
        TuskError: If ``output_backend`` names an uninstalled package.
    """
    keys = list(keys)
    if len(set(keys)) != len(keys):
        raise SchemaError(
            "X contains duplicate keys; the feature matrix is keyed by the "
            "target's primary key, so each row must be requested at most once",
        )

    frame = nw.from_native(lazy_native).lazy()
    filtered = frame.filter(nw.col(primary_key).is_in(keys))
    collected = _collect(filtered, output_backend)

    found = set(collected[primary_key].to_list())
    missing = [k for k in keys if k not in found]
    if missing:
        raise SchemaError(
            f"no row for {len(missing)} of {len(keys)} keys, e.g. "
            f"{missing[:5]}; they are absent from {primary_key!r} or were "
            f"excluded by cutoff_time",
        )

    # The key dtype must be taken from the collected frame rather than
    # inferred: inference gives int64, and a duckdb database keyed by int32
    # then fails the join with ArrowInvalid.
    order = nw.from_dict(
        {primary_key: keys, _POSITION: list(range(len(keys)))},
        schema={primary_key: collected.schema[primary_key], _POSITION: nw.Int64()},
        backend=collected.implementation,
    )
    joined = order.join(collected, on=primary_key, how="left").sort(_POSITION)
    return joined.drop([primary_key, _POSITION]).to_native()


def _collect(frame: nw.LazyFrame, output_backend: str | None) -> nw.DataFrame:
    """Collect, translating a missing backend package into a tusk error.

    Args:
        frame: The filtered lazy matrix.
        output_backend: Backend name, or None for the database's own.

    Returns:
        The collected frame.

    Raises:
        TuskError: If ``output_backend`` names a package that is not installed.
    """
    from tusk.exceptions import TuskError

    if output_backend is None:
        return frame.collect()
    try:
        return frame.collect(backend=output_backend)
    except ModuleNotFoundError as exc:
        raise TuskError(
            f"output_backend={output_backend!r} needs the {output_backend} "
            f"package, which is not installed; `uv add {output_backend}`",
        ) from exc


@contextlib.contextmanager
def backend_hint(frame: Any) -> Iterator[None]:
    """Attach a backend hint to whatever the user's pipeline raises.

    Many sklearn transformers reject non-pandas frames in ways that read as
    unrelated type errors -- ``ColumnTransformer`` on pyarrow gives
    ``TypeError: Index must either be string or integer``. The hint names the
    backend and the fix.

    The exception is re-raised **unchanged**. Wrapping it would change its
    type and break a user's ``except ValueError`` around their own pipeline.

    Args:
        frame: The matrix handed to the pipeline, used to name the backend.

    Yields:
        None.
    """
    try:
        yield
    except Exception as exc:
        backend = nw.from_native(frame, eager_only=True, pass_through=True)
        name = getattr(getattr(backend, "implementation", None), "name", "unknown")
        hint = (
            f"the feature matrix was collected as {name}; if this pipeline "
            f'needs pandas, set output_backend="pandas"'
        )
        if hasattr(exc, "add_note"):
            exc.add_note(hint)
        else:
            warnings.warn(hint, UserWarning, stacklevel=2)
        raise

import pytest

from tusk.sklearn._lineage import make_sentinels

COLUMNS = ["age", "MODE__transactions__category", "COUNT__transactions"]


def test_every_column_gets_a_distinct_sentinel():
    sentinels = make_sentinels(COLUMNS)
    assert len(set(sentinels.mapping.values())) == len(COLUMNS)


def test_sentinels_are_fixed_width_so_prefixes_cannot_collide():
    sentinels = make_sentinels([f"c{i}" for i in range(11)])
    lengths = {len(name) for name in sentinels.mapping.values()}
    assert len(lengths) == 1


def test_a_one_hot_name_resolves_to_its_source_column():
    sentinels = make_sentinels(COLUMNS)
    encoded = f"oh__{sentinels.mapping['MODE__transactions__category']}_a"
    assert sentinels.sources(encoded) == ["MODE__transactions__category"]


def test_a_multi_input_name_resolves_to_every_source():
    sentinels = make_sentinels(COLUMNS)
    encoded = f"{sentinels.mapping['age']} {sentinels.mapping['COUNT__transactions']}"
    assert sentinels.sources(encoded) == ["age", "COUNT__transactions"]


def test_an_opaque_name_resolves_to_nothing():
    assert make_sentinels(COLUMNS).sources("pca0") == []


def test_restore_puts_real_names_back():
    sentinels = make_sentinels(COLUMNS)
    encoded = f"oh__{sentinels.mapping['MODE__transactions__category']}_a"
    assert sentinels.restore(encoded) == "oh__MODE__transactions__category_a"

"""Reconnecting encoded columns to the tusk features they came from.

featuretools can prune feature definitions because ``encode_features`` returns
definitions for the one-hot columns, keeping the selector's mask and the
feature list 1:1. Handing encoding to scikit-learn severs that link: a support
mask indexes *encoded* columns and says nothing about which feature produced
them.

This module rebuilds the link by exploiting scikit-learn's own convention that
output names derive from input names. Rename the matrix columns to opaque
tokens, and every downstream name carries its provenance in plain sight.
"""

from __future__ import annotations

import re
import secrets
from collections.abc import Sequence
from dataclasses import dataclass, field


@dataclass(frozen=True)
class Sentinels:
    """An opaque renaming of matrix columns, and the map back.

    Attributes:
        prefix: Random per-fit token prefix.
        width: Zero-padded index width, fixed so no sentinel is a prefix of
            another.
        columns: The tusk-space column names, positionally indexed.
        mapping: Column name to sentinel.
    """

    prefix: str
    width: int
    columns: tuple[str, ...]
    mapping: dict[str, str] = field(compare=False)

    @property
    def _pattern(self) -> re.Pattern[str]:
        """Matches one sentinel and captures its index."""
        return re.compile(rf"{re.escape(self.prefix)}(\d{{{self.width}}})")

    def sources(self, name: str) -> list[str]:
        """Columns an encoded name derives from.

        Every match is collected rather than the first, so a multi-input
        transformer such as ``PolynomialFeatures`` correctly reports both of
        its sources and keeps both features alive.

        Args:
            name: One name from the encoder's ``get_feature_names_out()``.

        Returns:
            The tusk-space columns it mentions, in column order. Empty means
            lineage is unrecoverable for this name -- the caller must then keep
            every feature rather than guess.
        """
        found = {int(m) for m in self._pattern.findall(name)}
        return [self.columns[i] for i in sorted(found) if i < len(self.columns)]

    def restore(self, name: str) -> str:
        """Substitute sentinels back to real column names.

        Args:
            name: One name from the encoder's ``get_feature_names_out()``.

        Returns:
            The same name with every sentinel replaced, so users read
            ``oh__MODE__transactions__category_a`` rather than
            ``oh___t9f3a_0001_a``.
        """

        def replace(match: re.Match[str]) -> str:
            index = int(match.group(1))
            return self.columns[index] if index < len(self.columns) else match.group(0)

        return self._pattern.sub(replace, name)


def make_sentinels(columns: Sequence[str]) -> Sentinels:
    """Build an opaque renaming for one fit.

    The prefix is random per fit. A data *value* that happens to look like a
    sentinel would otherwise be misread as a column reference; randomness makes
    that vanishingly unlikely, and its only consequence is an extra source,
    which over-keeps rather than over-prunes.

    Args:
        columns: The matrix's column names, in order.

    Returns:
        The sentinel mapping and the map back.
    """
    prefix = f"_t{secrets.token_hex(4)}_"
    width = max(4, len(str(max(len(columns) - 1, 0))))
    mapping = {c: f"{prefix}{i:0{width}d}" for i, c in enumerate(columns)}
    return Sentinels(prefix, width, tuple(columns), mapping)

def test_categorical_family_matches_categorical_and_enum():
    import polars as pl

    frame = nw.from_native(
        pl.DataFrame(
            {
                "s": ["t"],
                "c": pl.Series(["x"]).cast(pl.Categorical),
                "e": pl.Series(["a"]).cast(pl.Enum(["a", "b"])),
            },
        ),
        eager_only=True,
    )
    schema = frame.schema
    assert matches(schema["c"], DtypeFamily.CATEGORICAL)
    assert matches(schema["e"], DtypeFamily.CATEGORICAL)
    assert not matches(schema["s"], DtypeFamily.CATEGORICAL)


def test_string_family_still_excludes_categorical_and_enum():
    import polars as pl

    frame = nw.from_native(
        pl.DataFrame(
            {
                "s": ["t"],
                "c": pl.Series(["x"]).cast(pl.Categorical),
                "e": pl.Series(["a"]).cast(pl.Enum(["a", "b"])),
            },
        ),
        eager_only=True,
    )
    schema = frame.schema
    assert matches(schema["s"], DtypeFamily.STRING)
    assert not matches(schema["c"], DtypeFamily.STRING)
    assert not matches(schema["e"], DtypeFamily.STRING)

    CATEGORICAL = "categorical"

    if family is DtypeFamily.CATEGORICAL:
        return dtype == nw.Categorical or dtype == nw.Enum

import datetime as dt

import polars as pl
import pytest
from sklearn.base import BaseEstimator, TransformerMixin
from sklearn.compose import ColumnTransformer
from sklearn.feature_selection import SelectKBest, f_classif
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import OneHotEncoder, StandardScaler

from tusk.dtypes import DtypeFamily
from tusk.exceptions import EncoderError
from tusk.sklearn import dtype_selector
from tusk.sklearn._encoders import encoder_prefix, selector_of, validate_inner

FRAME = pl.DataFrame(
    {"age": [20.0, 30.0], "cat": ["a", "b"], "cnt": [1.0, 2.0]},
)


def test_dtype_selector_splits_numeric_from_string():
    assert dtype_selector("numeric")(FRAME) == ["age", "cnt"]
    assert dtype_selector("string")(FRAME) == ["cat"]


def test_dtype_selector_keeps_booleans_and_dates_out_of_string():
    frame = pl.DataFrame(
        {
            "n": [1.0],
            "s": ["a"],
            "b": [True],
            "d": [dt.datetime(2024, 1, 1)],
        },
    )
    assert dtype_selector("numeric")(frame) == ["n"]
    assert dtype_selector("string")(frame) == ["s"]
    assert dtype_selector("boolean")(frame) == ["b"]
    assert dtype_selector("temporal")(frame) == ["d"]


def test_dtype_selector_separates_categorical_from_string():
    frame = pl.DataFrame(
        {
            "s": ["t"],
            "c": pl.Series(["x"]).cast(pl.Categorical),
            "e": pl.Series(["a"]).cast(pl.Enum(["a", "b"])),
        },
    )
    assert dtype_selector("string")(frame) == ["s"]
    assert dtype_selector(DtypeFamily.CATEGORICAL)(frame) == ["c", "e"]


def test_dtype_selector_rejects_an_unknown_family():
    with pytest.raises(ValueError):
        dtype_selector("texty")


def test_dtype_selector_reevaluates_on_a_subset():
    assert dtype_selector("numeric")(FRAME.select(["age"])) == ["age"]


def test_dtype_selector_works_inside_a_column_transformer_on_polars():
    encoder = ColumnTransformer(
        [
            ("oh", OneHotEncoder(handle_unknown="ignore"), dtype_selector("string")),
            ("num", StandardScaler(), dtype_selector("numeric")),
        ],
    ).fit(FRAME)
    assert list(encoder.get_feature_names_out()) == [
        "oh__cat_a",
        "oh__cat_b",
        "num__age",
        "num__cnt",
    ]


def test_selector_of_finds_the_last_step():
    selector = SelectKBest(f_classif, k=1)
    assert selector_of(Pipeline([("s", StandardScaler()), ("sel", selector)])) is selector
    assert selector_of(selector) is selector


def test_encoder_prefix_is_identity_when_there_is_no_encoder():
    prefix = encoder_prefix(SelectKBest(f_classif, k=1))
    fitted = prefix.fit(FRAME)
    assert list(fitted.get_feature_names_out()) == ["age", "cat", "cnt"]


def test_encoder_prefix_of_a_single_step_pipeline_is_identity():
    prefix = encoder_prefix(Pipeline([("sel", SelectKBest(f_classif, k=1))]))
    assert list(prefix.fit(FRAME).get_feature_names_out()) == ["age", "cat", "cnt"]


def test_inner_must_end_in_a_selector():
    with pytest.raises(EncoderError, match="SelectorMixin"):
        validate_inner(Pipeline([("s", StandardScaler())]))


def test_explicit_column_lists_are_refused():
    inner = Pipeline(
        [
            ("enc", ColumnTransformer([("num", StandardScaler(), ["age", "cnt"])])),
            ("sel", SelectKBest(f_classif, k=1)),
        ],
    )
    with pytest.raises(EncoderError, match="dtype_selector"):
        validate_inner(inner)


class _Opaque(BaseEstimator, TransformerMixin):
    """A transformer with no get_feature_names_out.

    TransformerMixin does not supply one -- only OneToOneFeatureMixin and
    ClassNamePrefixFeaturesOutMixin do -- so this needs no trickery. It stands
    in for scikit-lego's TypeSelector without taking the dependency.
    """

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

    def transform(self, X):
        return X


def test_a_step_without_feature_names_is_refused():
    inner = Pipeline(
        [("opaque", _Opaque()), ("sel", SelectKBest(f_classif, k=1))],
    )
    with pytest.raises(EncoderError, match="get_feature_names_out"):
        validate_inner(inner)


def test_a_callable_column_list_is_accepted():
    inner = Pipeline(
        [
            ("enc", ColumnTransformer([("num", StandardScaler(), dtype_selector("numeric"))])),
            ("sel", SelectKBest(f_classif, k=1)),
        ],
    )
    validate_inner(inner)

"""What the user's ``inner`` pipeline must look like, and why.

Pruning narrows the matrix, so the encoder is refit on fewer columns than it
first saw. That makes indifference to which columns are present a hard
requirement rather than a nicety, and it is the reason a whole shape of
``ColumnTransformer`` is refused here.
"""

from __future__ import annotations

from typing import Any

import narwhals as nw
from sklearn.compose import ColumnTransformer
from sklearn.feature_selection import SelectorMixin
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import FunctionTransformer

from tusk.dtypes import DtypeFamily, matches
from tusk.exceptions import EncoderError


class dtype_selector:  # noqa: N801
    """Select columns by :class:`~tusk.dtypes.DtypeFamily`, on any backend.

    scikit-learn's ``make_column_selector`` is the right shape but rejects
    anything that is not pandas, which collides with collecting natively. skrub
    and scikit-lego were surveyed too: skrub's selectors are not callable and
    ``ColumnTransformer`` rejects them, and scikit-lego's ``TypeSelector`` is a
    transformer with no ``get_feature_names_out``, which lineage depends on.
    So this is the same idea through narwhals.

    The families are tusk's own rather than a new set, so ``"string"`` means
    here exactly what it means to a primitive: ``String`` and not
    ``Categorical`` or ``Enum``.

    Being a callable rather than a fixed list is what makes the refit work: it
    re-evaluates against whatever frame it is handed, so narrowing the matrix
    narrows the selection instead of breaking it.

    Attributes:
        family: The :class:`~tusk.dtypes.DtypeFamily` to select.
    """

    def __init__(self, family: DtypeFamily | str) -> None:
        """Create a selector.

        Args:
            family: A ``DtypeFamily`` or its string value, such as
                ``"numeric"`` or ``"string"``.

        Raises:
            ValueError: If the string names no family. Raised by the enum, so
                the message lists the valid values.
        """
        self.family = DtypeFamily(family)

    def __call__(self, X: Any) -> list[str]:
        """Return the matching column names.

        Args:
            X: The frame the encoder is being fitted on.

        Returns:
            Matching column names, in frame order.
        """
        schema = nw.from_native(X, eager_only=True).schema
        return [c for c, d in schema.items() if matches(d, self.family)]

    def __repr__(self) -> str:
        """Show the family, so cloned estimators print readably."""
        return f"dtype_selector({self.family.value!r})"


def selector_of(inner: Any) -> Any:
    """The selector at the end of ``inner``.

    Args:
        inner: The user's encode-and-select estimator.

    Returns:
        Its last step if it is a pipeline, else ``inner`` itself.
    """
    return inner[-1] if isinstance(inner, Pipeline) else inner


def encoder_prefix(inner: Any) -> Any:
    """Everything in ``inner`` before the selector.

    A bare selector, or a one-step pipeline wrapping one, has no encoder; an
    identity transformer stands in so callers need no special case. Slicing a
    one-step pipeline would give an empty one, which cannot be fitted.

    Args:
        inner: The user's encode-and-select estimator.

    Returns:
        An unfitted estimator producing encoded space from tusk space.
    """
    if isinstance(inner, Pipeline) and len(inner.steps) > 1:
        return inner[:-1]
    return FunctionTransformer(feature_names_out="one-to-one")


def validate_inner(inner: Any) -> None:
    """Check ``inner`` against the contract, before anything is fitted.

    Args:
        inner: The user's encode-and-select estimator.

    Raises:
        EncoderError: If it does not end in a ``SelectorMixin``, or if any
            ``ColumnTransformer`` within it names its columns explicitly.
    """
    if not isinstance(selector_of(inner), SelectorMixin):
        raise EncoderError(
            "inner must end in a scikit-learn SelectorMixin -- its get_support() "
            "mask is what tells tusk which features to keep; got "
            f"{type(selector_of(inner)).__name__}",
        )
    _reject_explicit_columns(inner)
    _require_feature_names(encoder_prefix(inner))


def _require_feature_names(estimator: Any) -> None:
    """Refuse an encoder step that cannot report its output names.

    Lineage reads ``get_feature_names_out`` and nothing else, so a transformer
    without it is wholly opaque. Checking each step rather than the pipeline is
    the point: ``Pipeline`` always *has* the attribute and only fails when it
    reaches the step that does not, deep inside the fit and far from the cause.
    scikit-lego's ``TypeSelector`` is a real example.

    Args:
        estimator: An estimator to inspect, recursively.

    Raises:
        EncoderError: If any step cannot report output names.
    """
    if isinstance(estimator, Pipeline):
        for _, step in estimator.steps:
            _require_feature_names(step)
        return
    if isinstance(estimator, ColumnTransformer):
        for _, transformer, _columns in estimator.transformers:
            if transformer not in ("drop", "passthrough"):
                _require_feature_names(transformer)
        return
    if not hasattr(estimator, "get_feature_names_out"):
        raise EncoderError(
            f"{type(estimator).__name__} has no get_feature_names_out(), so "
            "tusk cannot tell which features its output columns came from. "
            "Every step before the selector must implement it.",
        )


def _reject_explicit_columns(estimator: Any) -> None:
    """Refuse a ``ColumnTransformer`` that names DFS-generated columns.

    Such an encoder is wrong here for three independent reasons: the names do
    not exist until DFS has run, so they must be hardcoded and go stale
    silently; ``remainder="drop"`` discards every generated feature the user
    did not name; and it is the one shape that cannot be refit on the pruned
    matrix.

    Args:
        estimator: An estimator to inspect, recursively.

    Raises:
        EncoderError: If an explicit column list is found.
    """
    if isinstance(estimator, Pipeline):
        for _, step in estimator.steps:
            _reject_explicit_columns(step)
    elif isinstance(estimator, ColumnTransformer):
        for name, transformer, columns in estimator.transformers:
            if not callable(columns):
                raise EncoderError(
                    f"ColumnTransformer step {name!r} names its columns "
                    f"explicitly ({list(columns)[:3]}...), which cannot be "
                    "refit once selection narrows the matrix, and DFS "
                    "generates its column names so they cannot be known in "
                    "advance anyway. Use a callable instead, such as "
                    "tusk.sklearn.dtype_selector('numeric') or "
                    "dtype_selector('string').",
                )
            _reject_explicit_columns(transformer)

from tusk.sklearn._encoders import dtype_selector

__all__ = ["dtype_selector"]

import warnings

import narwhals as nw
import numpy as np
import pytest
import sklearn
from sklearn.linear_model import LogisticRegression
from sklearn.model_selection import GridSearchCV, cross_val_score
from sklearn.pipeline import Pipeline

from tusk.exceptions import SchemaError
from tusk.sklearn import DFSTransformer

KEYS = [1, 2, 3]
Y = [0, 1, 0]


def _transformer():
    return DFSTransformer(target_table="customers", max_depth=2)


def test_fit_reads_no_rows_and_produces_feature_definitions(db):
    fitted = _transformer().fit(KEYS, database=db)
    assert fitted.features_
    assert all(f.table == "customers" for f in fitted.features_)


def test_transform_returns_one_row_per_key_in_key_order(db):
    fitted = _transformer().fit(KEYS, database=db)
    out = nw.from_native(fitted.transform([3, 1], database=db))
    assert out.shape[0] == 2


def test_the_primary_key_is_not_a_feature(db):
    fitted = _transformer().fit(KEYS, database=db)
    out = nw.from_native(fitted.transform(KEYS, database=db))
    assert "id" not in out.columns


def test_feature_names_match_the_matrix_columns(db):
    fitted = _transformer().fit(KEYS, database=db)
    out = nw.from_native(fitted.transform(KEYS, database=db))
    assert list(fitted.get_feature_names_out()) == list(out.columns)


def test_an_unknown_key_raises_rather_than_dropping_a_row(db):
    fitted = _transformer().fit(KEYS, database=db)
    with pytest.raises(SchemaError, match="no row"):
        fitted.transform([1, 999], database=db)


def test_transform_falls_back_to_the_database_seen_at_fit(db):
    fitted = _transformer().fit(KEYS, database=db)
    out = nw.from_native(fitted.transform(KEYS))
    assert out.shape[0] == 3


def test_it_routes_the_database_through_a_pipeline(db):
    with sklearn.config_context(enable_metadata_routing=True):
        pipe = Pipeline([("dfs", _transformer()), ("clf", LogisticRegression())])
        pipe.fit(np.array(KEYS).reshape(-1, 1), Y, database=db)
        assert len(pipe.predict(np.array([[1], [2]]), database=db)) == 2


def test_cross_val_score_produces_finite_scores(db):
    with sklearn.config_context(enable_metadata_routing=True):
        pipe = Pipeline([("dfs", _transformer()), ("clf", LogisticRegression())])
        scores = cross_val_score(
            pipe,
            np.array(KEYS).reshape(-1, 1),
            Y,
            cv=3,
            params={"database": db},
        )
    assert np.isfinite(scores).all()


def test_grid_search_can_tune_a_dfs_parameter(db):
    with sklearn.config_context(enable_metadata_routing=True):
        search = GridSearchCV(
            Pipeline([("dfs", _transformer()), ("clf", LogisticRegression())]),
            {"dfs__max_depth": [1, 2]},
            cv=3,
        )
        search.fit(np.array(KEYS).reshape(-1, 1), Y, database=db)
    assert search.best_params_["dfs__max_depth"] in (1, 2)


def test_fit_transform_forwards_the_database_without_warning(db):
    # TransformerMixin's fit_transform drops metadata on the way to transform,
    # and sklearn warns about exactly that. The database_ fallback would hide a
    # broken override, so the absence of the warning is what we assert.
    with sklearn.config_context(enable_metadata_routing=True):
        with warnings.catch_warnings():
            warnings.simplefilter("error", UserWarning)
            _transformer().fit_transform(KEYS, Y, database=db)


def test_importing_tusk_does_not_import_sklearn():
    import subprocess
    import sys

    result = subprocess.run(
        [sys.executable, "-c", "import sys, tusk; print('sklearn' in sys.modules)"],
        capture_output=True,
        text=True,
        check=True,
    )
    assert result.stdout.strip() == "False"

"""Deep feature synthesis as scikit-learn estimators.

``X`` is the target table's primary-key column, not the database. That
inversion is what makes ``X`` array-like, so scikit-learn can split it: the
keys *are* the rows, which makes misalignment against ``y`` impossible rather
than merely discouraged, and lets ``cross_val_score`` and ``GridSearchCV``
work at all. The database travels as routed metadata, because ``clone``
deep-copies constructor parameters and a duckdb relation cannot be pickled.
"""

from __future__ import annotations

from collections.abc import Iterable
from datetime import datetime
from typing import Any

import numpy as np
from sklearn.base import BaseEstimator, TransformerMixin
from sklearn.utils.validation import check_is_fitted

from tusk.api import apply_features
from tusk.database import Database
from tusk.exceptions import SchemaError
from tusk.primitives.aggregation import AGG_DEFAULTS
from tusk.primitives.base import Primitive
from tusk.primitives.transform import TRANS_DEFAULTS
from tusk.sklearn._frames import as_keys, collect_matrix
from tusk.synthesis import synthesize


class DFSTransformer(BaseEstimator, TransformerMixin):
    """Deep feature synthesis as a pipeline step.

    Attributes:
        features_: Feature definitions, set by :meth:`fit`.
        database_: The database seen at fit, used when none is routed later.
    """

    __metadata_request__fit = {"database": True}
    __metadata_request__transform = {"database": True}

    def __init__(
        self,
        target_table: str,
        agg_primitives: Iterable[str | Primitive] | None = None,
        trans_primitives: Iterable[str | Primitive] | None = None,
        groupby_trans_primitives: Iterable[str | Primitive] | None = None,
        max_depth: int = 2,
        cutoff_time: datetime | None = None,
        output_backend: str | None = None,
    ) -> None:
        """Configure synthesis.

        Every argument is stored unmodified under its own name: scikit-learn
        requires ``__init__`` to only assign, or ``get_params`` and ``clone``
        cannot round-trip the estimator.

        Args:
            target_table: Table to build features for.
            agg_primitives: Aggregation primitives; None selects the defaults.
            trans_primitives: Transform primitives; None selects the defaults.
            groupby_trans_primitives: Transforms within foreign-key groups.
            max_depth: Maximum stacked primitive applications.
            cutoff_time: Only rows at or before this are visible. Clone-safe,
                unlike the database, so it lives here and can be searched.
            output_backend: Backend to collect to. None collects natively,
                which keeps narwhals-native transformers on their own frame
                type.
        """
        self.target_table = target_table
        self.agg_primitives = agg_primitives
        self.trans_primitives = trans_primitives
        self.groupby_trans_primitives = groupby_trans_primitives
        self.max_depth = max_depth
        self.cutoff_time = cutoff_time
        self.output_backend = output_backend

    def fit(
        self,
        X: Any,
        y: Any = None,
        database: Database | None = None,
    ) -> DFSTransformer:
        """Synthesize feature definitions.

        Reads the schema and no rows, so this is cheap; the expensive pass is
        in :meth:`transform`.

        Args:
            X: The target's primary-key column. Unused here beyond validation,
                since synthesis depends only on the schema.
            y: Ignored; present for the scikit-learn signature.
            database: The database, routed as metadata.

        Returns:
            This estimator.

        Raises:
            SchemaError: If no database was supplied.
        """
        if database is None:
            raise SchemaError(
                "no database supplied; pass database=… to fit, and enable "
                "sklearn.set_config(enable_metadata_routing=True) when fitting "
                "inside a Pipeline",
            )
        as_keys(X)
        self.database_ = database
        self.features_ = list(
            synthesize(
                database=database,
                target_table=self.target_table,
                agg_primitives=AGG_DEFAULTS
                if self.agg_primitives is None
                else self.agg_primitives,
                trans_primitives=TRANS_DEFAULTS
                if self.trans_primitives is None
                else self.trans_primitives,
                groupby_trans_primitives=self.groupby_trans_primitives or (),
                max_depth=self.max_depth,
            ),
        )
        return self

    def transform(self, X: Any, database: Database | None = None) -> Any:
        """Compute the feature matrix for the keys in ``X``.

        Args:
            X: The target's primary-key column. Its order becomes the matrix's
                row order.
            database: The database, routed as metadata. When absent, the one
                seen at fit is used -- scikit-learn's scorers call ``predict``
                with no metadata, so without this fallback every
                cross-validated score would be ``nan``.

        Returns:
            An eager native frame, one row per key, in key order.
        """
        check_is_fitted(self, "features_")
        db = self.database_ if database is None else database
        return collect_matrix(
            apply_features(self.features_, db, self.cutoff_time),
            self._primary_key(db),
            as_keys(X),
            self.output_backend,
        )

    def fit_transform(
        self,
        X: Any,
        y: Any = None,
        database: Database | None = None,
        **kwargs: Any,
    ) -> Any:
        """Fit, then transform, forwarding the database to both.

        ``TransformerMixin.fit_transform`` drops metadata on the way to
        ``transform``; scikit-learn warns about exactly this. Overriding is the
        prescribed fix.

        Args:
            X: The target's primary-key column.
            y: Ignored.
            database: The database, routed as metadata.
            **kwargs: Ignored; absorbs scikit-learn's fit parameters.

        Returns:
            The feature matrix.
        """
        return self.fit(X, y, database=database).transform(X, database=database)

    def get_feature_names_out(self, input_features: Any = None) -> np.ndarray:
        """Column names of the matrix, in column order.

        A multi-output primitive contributes several names, so this is wider
        than ``len(features_)``.

        Args:
            input_features: Ignored; present for the scikit-learn signature.

        Returns:
            The names, as an object array.
        """
        check_is_fitted(self, "features_")
        names = [n for f in self.features_ for n in f.output_names]
        return np.asarray(names, dtype=object)

    def _primary_key(self, database: Database) -> str:
        """The target table's key column.

        Args:
            database: The database to read the schema from.

        Returns:
            The primary key column name.

        Raises:
            SchemaError: If the target table declares none.
        """
        primary_key = database.schema(self.target_table).primary_key
        if primary_key is None:
            raise SchemaError(
                f"target table {self.target_table!r} needs a primary_key: "
                "it is what X names",
            )
        return primary_key

from tusk.sklearn._encoders import dtype_selector
from tusk.sklearn._transformers import DFSTransformer

__all__ = ["DFSTransformer", "dtype_selector"]

import narwhals as nw
import numpy as np
import pytest
import sklearn
from sklearn.compose import ColumnTransformer
from sklearn.decomposition import PCA
from sklearn.feature_selection import SelectKBest, f_classif
from sklearn.linear_model import LogisticRegression
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import OneHotEncoder, StandardScaler

from tusk.exceptions import LineageWarning, UnencodedFeatureWarning
from tusk.sklearn import DFSSelectorTransformer, dtype_selector

KEYS = [1, 2, 3]
Y = [0, 1, 0]


def _inner(k=2):
    return Pipeline(
        [
            (
                "enc",
                ColumnTransformer(
                    [
                        (
                            "oh",
                            OneHotEncoder(handle_unknown="ignore"),
                            dtype_selector("string"),
                        ),
                        ("num", StandardScaler(), dtype_selector("numeric")),
                    ],
                ),
            ),
            ("sel", SelectKBest(f_classif, k=k)),
        ],
    )


def _transformer(k=2):
    return DFSSelectorTransformer(target_table="customers", inner=_inner(k), max_depth=2)


def test_fitting_prunes_the_feature_list(db):
    unpruned = len(_transformer().fit(KEYS, database=db).features_)
    from tusk.sklearn import DFSTransformer

    everything = len(
        DFSTransformer(target_table="customers", max_depth=2)
        .fit(KEYS, database=db)
        .features_,
    )
    assert unpruned < everything


def test_transform_width_equals_the_selectors_choice(db):
    fitted = _transformer(k=2).fit(KEYS, database=db)
    out = np.asarray(fitted.transform(KEYS, database=db))
    assert out.shape == (3, 2)


def test_pruned_features_are_never_computed(db):
    from tusk.sklearn import DFSTransformer

    everything = DFSTransformer(target_table="customers", max_depth=2).fit(
        KEYS,
        database=db,
    )
    all_columns = set(everything.get_feature_names_out())

    fitted = _transformer().fit(KEYS, database=db)
    kept = {n for f in fitted.features_ for n in f.output_names}
    dropped = all_columns - kept
    assert dropped, "selection kept every feature, so this proves nothing"

    # The parent's transform computes the tusk-space matrix for whatever
    # features_ currently holds, which after fit is the pruned list.
    computed = set(
        nw.from_native(DFSTransformer.transform(fitted, KEYS, database=db)).columns,
    )
    assert computed == kept
    assert not computed & dropped


def test_feature_names_out_are_readable_not_sentinels(db):
    fitted = _transformer().fit(KEYS, database=db)
    names = list(fitted.get_feature_names_out())
    assert names
    assert all(fitted.sentinels_.prefix not in name for name in names)


def test_an_opaque_encoder_keeps_every_feature_and_warns(db):
    inner = Pipeline([("pca", PCA(n_components=2)), ("sel", SelectKBest(f_classif, k=1))])
    transformer = DFSSelectorTransformer(
        target_table="customers",
        inner=inner,
        agg_primitives=["count"],
        trans_primitives=[],
    )
    from tusk.sklearn import DFSTransformer

    everything = len(
        DFSTransformer(
            target_table="customers",
            agg_primitives=["count"],
            trans_primitives=[],
        )
        .fit(KEYS, database=db)
        .features_,
    )
    with pytest.warns(LineageWarning):
        fitted = transformer.fit(KEYS, database=db)
    assert len(fitted.features_) == everything


def test_a_partial_encoder_warns_about_features_it_never_saw(db):
    # Only categorical columns are encoded, and remainder defaults to "drop",
    # so every numeric feature silently feeds nothing and gets pruned.
    inner = Pipeline(
        [
            (
                "enc",
                ColumnTransformer(
                    [
                        (
                            "oh",
                            OneHotEncoder(handle_unknown="ignore"),
                            dtype_selector("string"),
                        ),
                    ],
                ),
            ),
            ("sel", SelectKBest(f_classif, k=1)),
        ],
    )
    transformer = DFSSelectorTransformer(target_table="customers", inner=inner)
    with pytest.warns(UnencodedFeatureWarning, match="fed no encoded column"):
        transformer.fit(KEYS, Y, database=db)


def test_it_works_as_a_pipeline_step(db):
    with sklearn.config_context(enable_metadata_routing=True):
        pipe = Pipeline([("dfs", _transformer()), ("clf", LogisticRegression())])
        pipe.fit(np.array(KEYS).reshape(-1, 1), Y, database=db)
        assert len(pipe.predict(np.array([[1], [2]]), database=db)) == 2


def test_predicting_against_another_database_uses_the_pruned_features(db):
    fitted = _transformer().fit(KEYS, database=db)
    before = list(fitted.features_)
    fitted.transform([1, 2], database=db)
    assert fitted.features_ == before