API: preprocessing.imputation

class GroupImputerApplier(BaseApplier):
    """Fill gaps with the training value of each row's group, then with the global value.

    Rows whose group was unseen in training, whose group key is null, or whose
    group had no observed value fall back to the column's global training
    value. The scoring batch's own values are never used to compute fills.
    Ordered list keys may first try saved prefix statistics when the configured
    fallback is hierarchical. Legacy string keys retain their original lookup.
    """

    @staticmethod
    def validate_fitted_state(raw: dict) -> dict:
        """Inspect this node's supported saved state without fitting or applying data."""
        return _group_state(raw)

    @staticmethod
    def resolve_fitted_config(raw: dict, state: dict) -> dict:
        """Bind inference configuration to this node's inspected fitted artifact."""
        return _imputation_config("GroupImputer", fitted_columns(raw, state), state)

    @apply_method
    def apply(self, X: Any, y: Any, params: dict[str, Any]) -> Any:  # pylint: disable=arguments-differ
        """Fill the learned columns on the active engine; ``y`` passes through."""
        if not params.get("columns"):
            return X if y is None else (X, y)
        _require_group_column(_native(X).columns, params["group_by"])
        input_data = (X, y) if y is not None else X
        return apply_dual_engine(
            input_data, params, {"polars": self._apply_polars, "pandas": self._apply_pandas}
        )

    @staticmethod
    def _apply_pandas(X: Any, y: Any, params: dict[str, Any]) -> tuple[Any, Any]:
        out = X.copy()
        numeric = params["strategy"] in _NUMERIC_STRATEGIES
        for column in params["columns"]:
            if column not in out.columns:
                continue
            series = out[column]
            group_fill = _pandas_fills(out, series, column, params)
            if numeric:
                if not pd.api.types.is_float_dtype(series) or is_decimal_series(series):
                    series = pd.to_numeric(series).astype("float64")
                group_fill = pd.to_numeric(group_fill).astype("float64")
            series = _extend_categories(series, group_fill)
            out[column] = series.where(series.notna(), group_fill.to_numpy())
        return out, y

    @staticmethod
    def _apply_polars(X: Any, y: Any, params: dict[str, Any]) -> tuple[Any, Any]:
        numeric = params["strategy"] in _NUMERIC_STRATEGIES
        if "group_key_version" in params:
            exprs = [
                _polars_hierarchy_fill(X, column, params, numeric)
                for column in params["columns"]
                if column in X.columns
            ]
            return X.with_columns(exprs), y
        keys = pl.col(params["group_by"])
        key_dtype = X.schema[params["group_by"]]
        if isinstance(key_dtype, (pl.Categorical, pl.Enum)):
            keys = keys.cast(pl.Utf8)
            key_dtype = pl.String
        exprs = [
            _polars_fill(X.schema[column], column, keys, key_dtype, params, numeric)
            for column in params["columns"]
            if column in X.columns
        ]
        return X.with_columns(exprs), y

@NodeRegistry.register(
    "GroupImputer",
    GroupImputerApplier,
    execution_capabilities=tuple(
        ExecutionCapability(
            engine,
            "apply",
            execution_kind,
            "preserve",
            "row",
            config_match=(("strategy", strategy),),
        )
        for engine, execution_kind in (("pandas", "python_batch"), ("polars", "local"))
        for strategy in ("mean", "most_frequent", "median")
    ),
)
@node_meta(
    id="GroupImputer",
    name="Group Imputer",
    category="Preprocessing",
    description=(
        "Fills missing values with the mean, median or most frequent value of each row's "
        "group (for example per industry), falling back to the overall value."
    ),
    params={"group_by": "", "strategy": "mean", "columns": [], "fallback": "global"},
    learns_from_data=True,
)
class GroupImputerCalculator(BaseCalculator):
    """Learn per-group and global fill values from training rows.

    ``group_by`` names one column or an ordered list of group columns. Lists
    support ``fallback="hierarchical"`` to try successively shorter prefixes
    before the global training value; the default is ``"global"``. ``strategy`` is
    ``mean``, ``median`` or ``most_frequent`` (``mode`` is an alias). Leave
    ``columns`` empty to fill every numeric column (mean/median) or every
    column (most_frequent), excluding the group key and the target.
    """

    def infer_output_schema(
        self, input_schema: SkyulfSchema, config: dict[str, Any]
    ) -> SkyulfSchema | None:
        """Describe explicit fills; automatic numeric selection depends on row values."""
        if config.get("strategy", "mean") not in _NUMERIC_STRATEGIES:
            return input_schema
        if not config.get("columns"):
            # Automatic selection excludes constant and binary-valued columns,
            # which cannot be identified from dtypes alone.
            return None
        return promote_configured_columns_to_float64(input_schema, config)

    @fit_method
    def fit(self, X: Any, _y: Any, config: dict[str, Any]) -> GroupImputerArtifact:  # pylint: disable=arguments-differ
        """Compute ``[group, value]`` pairs and a global fallback for every filled column."""
        group_by = config.get("group_by")
        if not group_by:
            raise ValueError(
                "GroupImputer needs group_by: the column whose groups get their own fill value."
            )
        strategy = _strategy(config)
        fallback = group_fallback(config.get("fallback", "global"))
        if isinstance(group_by, list):
            group_by = group_columns(group_by)
        elif fallback != "global":
            raise ValueError(
                "GroupImputer hierarchical fallback requires an ordered group_by list."
            )
        X = _native(X)
        _require_group_column(X.columns, group_by)
        columns = _fill_columns(X, config, group_by, strategy)
        frame = _training_frame(X, columns, group_by, strategy)
        if isinstance(group_by, list):
            return fit_group_hierarchy(frame, config, columns, strategy)
        fill_values = {}
        for column in columns:
            values = frame[column]
            if strategy in _NUMERIC_STRATEGIES:
                values = pd.to_numeric(values)
            fill_values[column] = _statistic(values, strategy)
        return {
            "type": "group_imputer",
            "group_by": group_by,
            "strategy": strategy,
            "columns": columns,
            "group_values": {c: _group_values(frame, c, group_by, strategy) for c in columns},
            "fill_values": fill_values,
        }

class IterativeImputerApplier(BaseApplier):
    """Fill missing values from the fitted MICE (chained-equations) imputer."""

    @staticmethod
    def validate_inference_state(raw: dict) -> dict:
        """Inspect learned initial values and estimator sequence without executing them."""
        fields = {"type", "imputer_object", "columns", "estimator"}
        if local_state_fields(raw, "iterative_imputer", fields, allow_empty=True):
            imputer = _validate_local_imputer(raw, IterativeImputer)
            width = len(raw["columns"])
            _validate_iterative_initial(imputer, width)
            _validate_iterative_sequence(imputer, width)
        return raw

    @staticmethod
    def inference_capability(state: dict, *, engine: str) -> ExecutionCapability | None:
        """Admit deterministic saved rounds while abstaining from stochastic replay."""
        if engine not in {"pandas", "polars"}:
            return None
        IterativeImputerApplier.validate_inference_state(state)
        if state:
            imputer = state["imputer_object"]
            if not _deterministic_row_estimators(imputer):
                return None
        return ExecutionCapability(engine, "apply", "local", "preserve", "row")

    @apply_method
    def apply(self, X: Any, _y: Any, params: dict[str, Any]) -> Any:  # pylint: disable=arguments-differ
        """Impute ``X`` with the stored sklearn imputer; ``y`` passes through."""
        return apply_dual_engine(
            X, params, {"polars": self._apply_polars, "pandas": self._apply_pandas}
        )

    @staticmethod
    def _apply_polars(X: Any, _y: Any, params: dict[str, Any]) -> tuple[Any, Any]:
        cols = params.get("columns", [])
        imputer = params.get("imputer_object")
        if not resolve_valid_columns(X, cols) or not imputer:
            return X, _y
        return _sklearn_transform_subset(
            X, cols, imputer, is_polars=True, transform=_transform_iterative
        ), _y

    @staticmethod
    def _apply_pandas(X: Any, _y: Any, params: dict[str, Any]) -> tuple[Any, Any]:
        cols = params.get("columns", [])
        imputer = params.get("imputer_object")
        if not resolve_valid_columns(X, cols) or not imputer:
            return X, _y
        return _sklearn_transform_subset(
            X, cols, imputer, is_polars=False, transform=_transform_iterative
        ), _y

@NodeRegistry.register("IterativeImputer", IterativeImputerApplier)
@node_meta(
    id="IterativeImputer",
    name="Iterative Imputer (MICE)",
    category="Preprocessing",
    description="Multivariate imputation using chained equations.",
    params={"max_iter": 10, "random_state": 0, "estimator": "BayesianRidge", "columns": []},
    learns_from_data=True,
)
class IterativeImputerCalculator(BaseCalculator):
    """Fit a sklearn ``IterativeImputer`` with the configured estimator."""

    def infer_output_schema(
        self, input_schema: SkyulfSchema, config: dict[str, Any]
    ) -> SkyulfSchema:
        """Return a schema with selected columns promoted to ``float64``."""
        return promote_configured_columns_to_float64(input_schema, config)

    @fit_method
    def fit(self, X: Any, _y: Any, config: dict[str, Any]) -> IterativeImputerArtifact:  # pylint: disable=arguments-differ
        """Fit chained-equations imputation, dropping all-missing columns sklearn cannot impute."""
        estimator_name = config.get("estimator", "BayesianRidge")
        estimator = _build_iterative_estimator(estimator_name)
        if user_picked_no_columns(config):
            return {}

        max_iter = config.get("max_iter", 10)
        random_state = config.get("random_state", 0)

        # KNN/Iterative imputers always fit through numpy — engine choice
        # doesn't affect the fit math, so we skip the Pandas hop entirely.
        X_np, cols = resolve_columns_then_to_numpy(X, config, detect_numeric_columns)
        if not cols:
            return {}

        # sklearn silently drops all-missing columns from transform() output,
        # which desyncs the artifact's column list from the imputer's width.
        X_np, cols = drop_all_missing_columns(X_np, cols, "IterativeImputerCalculator")
        if not cols:
            return {}

        imputer = IterativeImputer(
            estimator=estimator, max_iter=max_iter, random_state=random_state
        )
        imputer.fit(X_np)

        return {
            "type": "iterative_imputer",
            "imputer_object": imputer,  # Not JSON serializable
            "columns": cols,
            "estimator": estimator_name,
        }

class KNNImputerApplier(BaseApplier):
    """Fill missing values from the fitted k-nearest-neighbors imputer."""

    @staticmethod
    def validate_inference_state(raw: dict) -> dict:
        """Inspect the saved neighbor estimator without refitting or normalizing it."""
        fields = {"type", "imputer_object", "columns", "n_neighbors", "weights"}
        if local_state_fields(raw, "knn_imputer", fields, allow_empty=True):
            imputer = _validate_local_imputer(raw, KNNImputer)
            _validate_knn_neighbors(imputer, len(raw["columns"]))
        return raw

    @staticmethod
    def inference_capability(state: dict, *, engine: str) -> ExecutionCapability | None:
        """Describe native neighbor replay; arbitrary metrics or weights stay undeclared."""
        if engine not in {"pandas", "polars"}:
            return None
        KNNImputerApplier.validate_inference_state(state)
        if state:
            imputer = state["imputer_object"]
            if callable(imputer.weights) or callable(imputer.metric):
                return None
        return ExecutionCapability(engine, "apply", "local", "preserve", "row")

    @apply_method
    def apply(self, X: Any, _y: Any, params: dict[str, Any]) -> Any:  # pylint: disable=arguments-differ
        """Impute ``X`` with the stored sklearn imputer; ``y`` passes through."""
        return apply_dual_engine(
            X, params, {"polars": self._apply_polars, "pandas": self._apply_pandas}
        )

    @staticmethod
    def _apply_polars(X: Any, _y: Any, params: dict[str, Any]) -> tuple[Any, Any]:
        cols = params.get("columns", [])
        imputer = params.get("imputer_object")
        if not resolve_valid_columns(X, cols) or not imputer:
            return X, _y
        return _sklearn_transform_subset(X, cols, imputer, is_polars=True), _y

    @staticmethod
    def _apply_pandas(X: Any, _y: Any, params: dict[str, Any]) -> tuple[Any, Any]:
        cols = params.get("columns", [])
        imputer = params.get("imputer_object")
        if not resolve_valid_columns(X, cols) or not imputer:
            return X, _y
        return _sklearn_transform_subset(X, cols, imputer, is_polars=False), _y

@NodeRegistry.register("KNNImputer", KNNImputerApplier)
@node_meta(
    id="KNNImputer",
    name="KNN Imputer",
    category="Preprocessing",
    description="Impute missing values using k-Nearest Neighbors.",
    params={"n_neighbors": 5, "weights": "uniform", "columns": []},
    learns_from_data=True,
)
class KNNImputerCalculator(BaseCalculator):
    """Fit a sklearn ``KNNImputer`` on the selected numeric columns."""

    def infer_output_schema(
        self, input_schema: SkyulfSchema, config: dict[str, Any]
    ) -> SkyulfSchema:
        """Return a schema with selected columns promoted to ``float64``."""
        return promote_configured_columns_to_float64(input_schema, config)

    @fit_method
    def fit(self, X: Any, _y: Any, config: dict[str, Any]) -> KNNImputerArtifact:  # pylint: disable=arguments-differ
        """Fit the neighbor imputer, dropping all-missing columns sklearn cannot impute."""
        if user_picked_no_columns(config):
            return {}

        n_neighbors = config.get("n_neighbors", 5)
        weights = config.get("weights", "uniform")

        # KNN/Iterative imputers always fit through numpy — engine choice
        # doesn't affect the fit math, so we skip the Pandas hop entirely.
        X_np, cols = resolve_columns_then_to_numpy(X, config, detect_numeric_columns)
        if not cols:
            return {}

        # sklearn silently drops all-missing columns from transform() output,
        # which desyncs the artifact's column list from the imputer's width.
        X_np, cols = drop_all_missing_columns(X_np, cols, "KNNImputerCalculator")
        if not cols:
            return {}

        imputer = KNNImputer(n_neighbors=n_neighbors, weights=weights)
        imputer.fit(X_np)

        return {
            "type": "knn_imputer",
            "imputer_object": imputer,  # Not JSON serializable
            "columns": cols,
            "n_neighbors": n_neighbors,
            "weights": weights,
        }

class SimpleImputerApplier(BaseApplier):
    """Apply fitted Simple Imputer fill values to missing values in selected columns.

    The calculator artifact records per-column values for ``mean``, ``median``,
    ``most_frequent`` (also accepted as ``mode``), or ``constant`` strategies.
    Missing columns seen during fitting are restored with their stored value,
    preserving the row count even when no input columns remain.
    Empty saved artifacts declare identity replay only when the recipe explicitly
    selects no columns; automatic or missing selections remain unreviewed.
    Spark supports only ``mean`` and ``constant`` artifacts and applies them
    with native expressions. All-missing means leave existing columns untouched.
    """

    @staticmethod
    def validate_fitted_state(raw: dict) -> dict:
        """Inspect this node's supported saved state without fitting or applying data."""
        return (
            _local_statistic_state(raw)
            if raw.get("strategy") in ("median", "most_frequent")
            else validate_state("SimpleImputer", raw)
        )

    @staticmethod
    def resolve_fitted_config(raw: dict, state: dict) -> dict:
        """Bind inference configuration to this node's inspected fitted artifact."""
        if not state:
            if not user_picked_no_columns(raw) or raw.get("_auto_columns"):
                raise ValueError("An empty imputer artifact requires explicit empty columns.")
            strategy = raw.get("strategy")
            if strategy in ("median", "most_frequent", "mode"):
                canonical = "median" if strategy == "median" else "most_frequent"
                return _imputation_config(
                    "SimpleImputer",
                    fitted_columns(raw, {"columns": []}),
                    {"strategy": canonical},
                )
        if state.get("strategy") in ("median", "most_frequent"):
            return _imputation_config("SimpleImputer", fitted_columns(raw, state), state)
        return portable_config("SimpleImputer", raw, state)

    @apply_method
    def apply(self, X: Any, _y: Any, params: dict[str, Any]) -> Any:  # pylint: disable=arguments-differ
        """Fill nulls and NaNs with the fitted per-column values; ``y`` passes through."""
        return apply_dual_engine(
            X,
            params,
            {
                "polars": self._apply_polars,
                "pandas": self._apply_pandas,
                "spark": apply_spark_imputer,
            },
        )

    @staticmethod
    def _apply_polars(X: Any, _y: Any, params: dict[str, Any]) -> tuple[Any, Any]:
        cols = params.get("columns", [])
        fill_values = params.get("fill_values", {})
        if not cols:
            return X, _y

        exprs: list[Any] = []
        for col in X.columns:
            if col in cols and col in fill_values:
                val = fill_values[col]
                # An all-null fitted column has no computable mean/median, so
                # its stored fill value is None. Skip it (leave the column
                # untouched) to mirror the pandas branch's `if val is None`
                # guard; fill_null(None) would raise.
                if val is None:
                    exprs.append(pl.col(col))
                    continue
                expr = pl.col(col).fill_null(val)
                if X.schema[col].is_float():
                    expr = expr.fill_nan(val)
                exprs.append(expr.alias(col))
            else:
                exprs.append(pl.col(col))

        # Restore columns that were present at fit time but missing in input X.
        exprs.extend(
            pl.repeat(pl.lit(fill_values[col]), pl.len()).alias(col)
            for col in cols
            if col not in X.columns and col in fill_values
        )

        return X.select(exprs), _y

    @staticmethod
    def _apply_pandas(X: Any, _y: Any, params: dict[str, Any]) -> tuple[Any, Any]:
        cols = params.get("columns", [])
        fill_values = params.get("fill_values", {})
        if not cols:
            return X, _y

        X_out = X.copy()
        for col in cols:
            val = fill_values.get(col)
            if val is None:
                continue
            if col not in X_out.columns:
                X_out[col] = val
            else:
                X_out[col] = _fill_pandas_series(X_out[col], val, params.get("strategy"))
        return X_out, _y

@NodeRegistry.register(
    "SimpleImputer",
    SimpleImputerApplier,
    execution_capabilities=tuple(
        ExecutionCapability(
            "spark",
            operation,
            "native",
            "preserve",
            "global" if operation == "fit" else "row",
            codec_version=1,
            config_match=(("strategy", strategy),),
        )
        for operation in ("fit", "apply")
        for strategy in ("mean", "constant")
    )
    + tuple(
        ExecutionCapability(
            engine,
            "apply",
            execution_kind,
            "preserve",
            "row",
            codec_version=1,
            config_match=(("strategy", strategy),),
        )
        for engine, execution_kind in (("pandas", "python_batch"), ("polars", "local"))
        for strategy in ("mean", "constant")
    )
    + tuple(
        ExecutionCapability(
            engine,
            "apply",
            execution_kind,
            "preserve",
            "row",
            config_match=(("strategy", strategy),),
        )
        for engine, execution_kind in (("pandas", "python_batch"), ("polars", "local"))
        for strategy in ("most_frequent", "median")
    ),
)
@node_meta(
    id="SimpleImputer",
    name="Simple Imputer",
    category="Preprocessing",
    description="Imputes missing values using mean, median, or constant.",
    params={"strategy": "mean", "fill_value": None, "columns": []},
    learns_from_data=True,
)
class SimpleImputerCalculator(BaseCalculator):
    """Fit per-column values for filling missing data with sklearn-compatible strategies.

    Supported ``strategy`` values are ``mean``, ``median``, ``most_frequent``,
    and ``constant``; ``mode`` is normalized to ``most_frequent``. Use
    ``columns`` to select columns and ``fill_value`` with ``constant``.
    Mean and median auto-select numeric columns and reject explicitly selected
    non-numeric columns. The other strategies can operate on all selected columns.
    Spark supports only ``mean`` and ``constant`` through a native aggregate;
    string and boolean constants require explicit compatible fill values.
    """

    def infer_output_schema(
        self, input_schema: SkyulfSchema, config: dict[str, Any]
    ) -> SkyulfSchema | None:
        """Preserve known double targets and abstain when output needs values or engine identity."""
        strategy = config.get("strategy", "mean")
        columns = config.get("columns")
        if not isinstance(columns, list):
            return None
        selected = [column for column in columns if column in input_schema]
        if not selected:
            return input_schema
        if any(input_schema.dtypes.get(column, "").lower() != "float64" for column in selected):
            return None
        if strategy == "constant" and type(config.get("fill_value")) not in (
            type(None),
            int,
            float,
            str,
        ):
            return None
        return input_schema

    @fit_method
    def fit(self, X: Any, _y: Any, config: dict[str, Any]) -> SimpleImputerArtifact:  # pylint: disable=arguments-differ
        """Resolve strategy and columns, then fit fill values on the active engine.

        ``mode`` is accepted as an alias of ``most_frequent``; mean/median are
        restricted to numeric columns.
        """
        return cast(
            SimpleImputerArtifact,
            fit_dual_engine(
                X,
                config,
                {
                    "polars": self._fit_polars_input,
                    "pandas": self._fit_pandas_input,
                    "spark": fit_spark_imputer,
                },
            ),
        )

    @staticmethod
    def _prepare_local(X: Any, config: dict[str, Any]) -> dict[str, Any] | None:
        """Keep existing local column/strategy resolution separate from Spark preparation."""
        if user_picked_no_columns(config):
            return None

        strategy = config.get("strategy", "mean")
        if strategy == "mode":
            strategy = "most_frequent"
        fill_value = config.get("fill_value")

        cols = _resolve_simple_columns(X, config, strategy)
        if not cols:
            return None

        if strategy in {"mean", "median"}:
            numeric = set(auto_detect_numeric_columns(X))
            invalid = [col for col in cols if col not in numeric]
            if invalid:
                raise ValueError(
                    f"SimpleImputer strategy '{strategy}' requires numeric columns; "
                    f"non-numeric columns selected: {invalid}."
                )

        # Stash resolved-once values into params so dispatched fits don't redo work.
        merged = dict(config)
        merged["_resolved_strategy"] = strategy
        merged["_resolved_cols"] = cols
        merged["_resolved_fill_value"] = fill_value

        return merged

    @staticmethod
    def _fit_polars_input(X: Any, y: Any, config: dict[str, Any]) -> dict[str, Any]:
        """Prepare the existing Polars contract before its fit implementation."""
        prepared = SimpleImputerCalculator._prepare_local(X, config)
        return {} if prepared is None else SimpleImputerCalculator._fit_polars(X, y, prepared)

    @staticmethod
    def _fit_pandas_input(X: Any, y: Any, config: dict[str, Any]) -> dict[str, Any]:
        """Prepare the existing pandas contract before its fit implementation."""
        prepared = SimpleImputerCalculator._prepare_local(X, config)
        return {} if prepared is None else SimpleImputerCalculator._fit_pandas(X, y, prepared)

    @staticmethod
    def _fit_polars(X: Any, _y: Any, params: dict[str, Any]) -> dict[str, Any]:
        cols: list[str] = params["_resolved_cols"]
        strategy: str = params["_resolved_strategy"]
        fill_value = params["_resolved_fill_value"]

        fill_values = (
            constant_text_values(X, cols, fill_value)
            if strategy == "constant" and isinstance(fill_value, str)
            else _compute_polars_fill_values(X, cols, strategy, fill_value)
        )
        missing_counts, total_missing = _polars_missing_counts(X, cols)

        return {
            "type": "simple_imputer",
            "strategy": strategy,
            "fill_values": fill_values,
            "columns": cols,
            "missing_counts": missing_counts,
            "total_missing": total_missing,
        }

    @staticmethod
    def _fit_pandas(X: Any, _y: Any, params: dict[str, Any]) -> dict[str, Any]:
        cols: list[str] = params["_resolved_cols"]
        strategy: str = params["_resolved_strategy"]
        fill_value = params["_resolved_fill_value"]

        # Mean/median: extra safety filter to numeric columns only.
        if strategy in ("mean", "median"):
            # Automatic selection already applied its cardinality heuristics.
            # Do not discard explicitly selected binary or constant columns here.
            numeric = set(detect_numeric_columns(X, exclude_binary=False, exclude_constant=False))
            cols = [c for c in cols if c in numeric]
            if not cols:
                return {}

        fill_values = _pandas_fill_values(X, cols, strategy, fill_value)
        missing_counts = X[cols].isnull().sum().to_dict()
        total_missing = int(sum(missing_counts.values()))

        return {
            "type": "simple_imputer",
            "strategy": strategy,
            "fill_values": fill_values,
            "columns": cols,
            "missing_counts": missing_counts,
            "total_missing": total_missing,
        }