API: preprocessing.transformations

class GeneralTransformationApplier(BaseApplier):
    """Apply per-column simple ops or fitted Box-Cox/Yeo-Johnson transforms."""

    @staticmethod
    def validate_inference_state(raw: dict) -> dict:
        """Inspect ordered saved rules while retaining the existing simple-rule checks."""
        local_state_fields(raw, "general_transformation", {"type", "transformations"})
        rules = raw["transformations"]
        if type(rules) not in (list, tuple):
            raise ValueError("GeneralTransformation rules must be a list or tuple.")
        for rule in rules:
            _validate_general_rule(rule)
        return raw

    @staticmethod
    def inference_capability(state: dict, *, engine: str) -> ExecutionCapability | None:
        """Retain conservative whole-request scope for fitted power transforms."""
        if engine not in ("pandas", "polars"):
            return None
        GeneralTransformationApplier.validate_inference_state(state)
        has_power = any(rule.get("method") in _POWER_METHODS for rule in state["transformations"])
        return ExecutionCapability(
            engine, "apply", "local", "preserve", "global" if has_power else "row"
        )

    @apply_method
    def apply(self, X: Any, _y: Any, params: dict[str, Any]) -> Any:  # pylint: disable=arguments-differ
        """Run each configured transformation on the active engine; ``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]:
        transformations = params.get("transformations", [])
        if not transformations:
            return X, _y

        X_out = X
        for item in transformations:
            col = item.get("column")
            method = item.get("method")
            if col not in X_out.columns:
                continue
            if method in _POWER_METHODS:
                X_out = _apply_power_to_polars_col(X_out, item)
                continue
            op = _POLARS_OPS.get(method)
            if op is None:
                continue
            X_out = _apply_polars_op(X_out, item, op)
        return X_out, _y

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

        df_out = X.copy()
        for item in transformations:
            col = item.get("column")
            method = item.get("method")
            if col not in df_out.columns:
                continue
            if method in _POWER_METHODS:
                df_out = _apply_power_to_pandas_col(df_out, item)
                continue
            op = _PANDAS_OPS.get(method)
            if op is None:
                continue
            df_out[col] = op(pd.to_numeric(df_out[col], errors="coerce"), item)
        return df_out, _y

@NodeRegistry.register("GeneralTransformation", GeneralTransformationApplier)
@node_meta(
    id="GeneralTransformation",
    name="General Transformation",
    category="Preprocessing",
    description="Apply various function transformations (log, sqrt, square, exp) to columns.",
    params={"transformations": []},
    learns_from_data=True,
)
class GeneralTransformationCalculator(BaseCalculator):
    """Resolve configured transformations, fitting lambdas for the power methods.

    Rules run in list order. Each power rule is fitted on the training values
    produced by earlier rules on the same column, matching artifact replay.
    Each power rule accepts ``standardize`` (default ``True``). Set it to
    ``False`` to apply only the fitted Box-Cox or Yeo-Johnson transform.
    The artifact retains the choice for subsequent transforms.
    """

    def infer_output_schema(
        self, input_schema: SkyulfSchema, config: dict[str, Any]
    ) -> SkyulfSchema:
        """Return updated dtypes for transformed numeric columns."""
        schema = input_schema
        touched: set[str] = set()
        for item in config.get("transformations", []):
            col = item.get("column")
            method = item.get("method")
            if col not in schema.columns or method is None:
                continue
            if method in _FLOAT_METHODS:
                touched.add(col)

        for col in touched:
            schema = schema.with_dtype(col, "float64")
        return schema

    @fit_method
    def fit(self, X: Any, _y: Any, config: dict[str, Any]) -> GeneralTransformationArtifact:  # pylint: disable=arguments-differ
        """Fit power lambdas in rule order; retain simple ops in the artifact.

        Box-Cox is skipped (with a warning) for non-positive values at that
        point in the rule sequence. A fit failure logs and skips that rule.
        """
        # Config: {'transformations': [{'column': 'col1', 'method': 'log'},
        #                              {'column': 'col2', 'method': 'yeo-johnson'}]}
        is_polars = get_engine(X).name == EngineName.POLARS
        fitted_transformations: list[dict[str, Any]] = []
        transformations = config.get("transformations", [])
        last_power_rule = {
            item.get("column"): index
            for index, item in enumerate(transformations)
            if item.get("method") in _POWER_METHODS
        }
        working_X = X
        applier = GeneralTransformationApplier()

        for index, item in enumerate(transformations):
            col = item.get("column")
            if col not in X.columns:
                continue

            fitted_item = _fit_transformation_rule(working_X, col, item, is_polars)
            if fitted_item is None:
                continue

            fitted_transformations.append(fitted_item)
            # Only materialize rules whose output is needed by a later fit.
            # Reuse replay semantics without mutating the caller's training data.
            if index < last_power_rule.get(col, -1):
                # apply_method exposes the public (data, params) call signature.
                working_X = applier.apply(  # pylint: disable=no-value-for-parameter
                    working_X, {"transformations": [fitted_item]}
                )

        return {
            "type": "general_transformation",
            "transformations": fitted_transformations,
        }

class PowerTransformerApplier(BaseApplier):
    """Apply the fitted Box-Cox / Yeo-Johnson transform to the selected columns."""

    @staticmethod
    def validate_inference_state(raw: dict) -> dict:
        """Inspect the saved lambdas and scaler without normalizing or relearning them."""
        if local_state_fields(
            raw,
            "power_transformer",
            {"type", "columns", "method", "standardize", "lambdas", "scaler_params"},
            allow_empty=True,
        ):
            columns = raw["columns"]
            if type(columns) is not list or any(not isinstance(col, str) for col in columns):
                raise ValueError("Fitted power columns must be a list of strings.")
            _validate_finite_power_state(raw, len(columns))
        return raw

    @staticmethod
    def inference_capability(state: dict, *, engine: str) -> ExecutionCapability | None:
        """Declare row-local filtering of invalid selected values and transformed outputs."""
        if engine not in ("pandas", "polars"):
            return None
        PowerTransformerApplier.validate_inference_state(state)
        return ExecutionCapability(
            engine, "apply", "local", "filter" if state else "preserve", "row"
        )

    @apply_method
    def apply(self, X: Any, _y: Any, params: dict[str, Any]) -> Any:  # pylint: disable=arguments-differ
        """Transform valid rows with stored lambdas and align ``y`` by position.

        Invalid numeric rows are excluded and counted in a warning. Missing
        values remain missing for later imputers. Valid
        neighbours keep their fitted transform. Missing fitted columns are
        warned about; malformed state and unexpected implementation errors raise.
        """
        if params.get("lambdas") is not None:
            _validate_finite_power_state(params, len(params["columns"]))
        return apply_dual_engine(
            (X, _y) if _y is not None else 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", [])
        if params.get("lambdas") is None:
            return X, _y
        valid_cols = _fitted_columns_present(X, cols)
        if not valid_cols:
            return X, _y

        X_vals = np.asarray(X.select(valid_cols).to_numpy(), dtype=np.float64)
        X_trans = _power_transform_array(X_vals, params, cols, valid_cols)
        series = [pl.Series(name, X_trans[:, i]) for i, name in enumerate(valid_cols)]
        positions = _valid_power_positions(X_trans, X_vals)
        return X.with_columns(series)[positions], select_rows_by_position(_y, positions)

    @staticmethod
    def _apply_pandas(X: Any, _y: Any, params: dict[str, Any]) -> tuple[Any, Any]:
        cols = params.get("columns", [])
        if params.get("lambdas") is None:
            return X, _y
        valid_cols = _fitted_columns_present(X, cols)
        if not valid_cols:
            return X, _y

        subset = decimal_columns_to_float(X[valid_cols], valid_cols)
        X_vals = subset.to_numpy(dtype=np.float64, na_value=np.nan)
        X_trans = _power_transform_array(X_vals, params, cols, valid_cols)
        df_out = X.copy()
        # Replace selected columns so nullable integers and Decimal objects
        # can receive float results without incompatible-dtype assignment.
        df_out[valid_cols] = X_trans
        positions = _valid_power_positions(X_trans, X_vals)
        return df_out.iloc[positions], select_rows_by_position(_y, positions)

@NodeRegistry.register("PowerTransformer", PowerTransformerApplier)
@node_meta(
    id="PowerTransformer",
    name="Power Transformer",
    category="Preprocessing",
    description="Apply a power transform featurewise to make data more Gaussian-like.",
    params={"method": "yeo-johnson", "standardize": True, "columns": []},
    learns_from_data=True,
)
class PowerTransformerCalculator(BaseCalculator):
    """Fit one shared PowerTransformer across the selected columns."""

    def infer_output_schema(
        self, input_schema: SkyulfSchema, config: dict[str, Any]
    ) -> SkyulfSchema:
        """Return a schema with transformed 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]) -> PowerTransformerArtifact:  # pylint: disable=arguments-differ
        """Fit selected columns on valid rows, excluding invalid rows from learned statistics."""
        if user_picked_no_columns(config):
            return {}

        method = config.get("method", "yeo-johnson")
        standardize = config.get("standardize", True)
        X_pd, cols = resolve_columns_then_to_pandas(X, config, detect_numeric_columns)
        if not cols:
            return {}
        values = _valid_training_values(X_pd, cols, method)
        return _fit_power_state(values, cols, method, standardize, len(X_pd))

class SimpleTransformationApplier(BaseApplier):
    """Apply the configured per-column math transformations (log, sqrt, ...)."""

    @staticmethod
    def validate_inference_state(raw: dict) -> dict:
        """Inspect ordered native rules without consuming iterators or normalizing state."""
        local_state_fields(raw, "simple_transformation", {"type", "transformations"})
        transformations = raw["transformations"]
        if transformations is None:
            return raw
        if type(transformations) not in (list, tuple):
            raise ValueError("SimpleTransformation rules must be a list or tuple.")
        for item in transformations:
            _validate_simple_rule(item)
        return raw

    @staticmethod
    def inference_capability(state: dict, *, engine: str) -> ExecutionCapability | None:
        """Describe native pointwise math as local row-preserving inference."""
        if engine not in ("pandas", "polars"):
            return None
        SimpleTransformationApplier.validate_inference_state(state)
        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
        """Run each configured transformation on the active engine; ``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]:
        transformations = params.get("transformations", [])
        if not transformations:
            return X, _y

        X_out = X
        for item in transformations:
            col = item.get("column")
            method = item.get("method")
            if col not in X_out.columns:
                continue
            op = _POLARS_OPS.get(method)
            if op is None:
                continue
            X_out = _apply_polars_op(X_out, item, op)
        return X_out, _y

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

        df_out = X.copy()
        for item in transformations:
            col = item.get("column")
            method = item.get("method")
            if col not in df_out.columns:
                continue
            op = _PANDAS_OPS.get(method)
            if op is None:
                continue
            df_out[col] = op(pd.to_numeric(df_out[col], errors="coerce"), item)
        return df_out, _y

@NodeRegistry.register("SimpleTransformation", SimpleTransformationApplier)
@node_meta(
    id="SimpleTransformation",
    name="Simple Transformation",
    category="Preprocessing",
    description="Apply simple mathematical transformations (log, sqrt, etc.).",
    params={"transformations": []},
    learns_from_data=False,
)
class SimpleTransformationCalculator(BaseCalculator):
    """Package the configured transformations into an artifact without reading the data."""

    _FLOAT_METHODS = {
        "square",
        "log",
        "sqrt",
        "square_root",
        "cube_root",
        "reciprocal",
        "exp",
        "exponential",
    }

    def infer_output_schema(
        self, input_schema: SkyulfSchema, config: dict[str, Any]
    ) -> SkyulfSchema:
        """Return updated dtypes for transformed numeric columns."""
        schema = input_schema
        touched: set[str] = set()
        for item in config.get("transformations", []):
            col = item.get("column")
            method = item.get("method")
            if col not in schema.columns or method is None:
                continue
            if method in self._FLOAT_METHODS:
                touched.add(col)

        for col in touched:
            schema = schema.with_dtype(col, "float64")
        return schema

    def fit(
        self,
        df: pd.DataFrame | SkyulfDataFrame | tuple[Any, ...] | Any,
        config: dict[str, Any],
    ) -> SimpleTransformationArtifact:
        """Validate requested math methods and save the ordered rules without fitting."""
        # Config: {'transformations': [{'column': 'col1', 'method': 'log'}, ...]}
        transformations = config.get("transformations", [])
        artifact: SimpleTransformationArtifact = {
            "type": "simple_transformation",
            "transformations": transformations,
        }
        SimpleTransformationApplier.validate_inference_state(cast(dict[str, Any], artifact))
        for item in transformations or []:
            method = item.get("method")
            if not isinstance(method, str) or method not in _PANDAS_OPS:
                raise ValueError(
                    f"Unknown SimpleTransformation method {method!r}. "
                    f"Supported methods: {', '.join(_PANDAS_OPS)}."
                )
        return artifact