API: preprocessing.casting

class CastingApplier(BaseApplier):
    """Cast each column named in ``type_map`` to its target dtype on the active engine.

    ``coerce_on_error=True`` (the default) is best-effort per column: the pandas
    path swallows any cast failure and leaves that column untouched, the polars
    path casts with ``strict=False``. With it off, both engines raise. The
    awkward families are hand-written on each path because the native casts
    disagree — polars truncates fractional floats where pandas masks them, and
    pandas' ``astype`` clamps out-of-range integers where polars nulls or raises.
    Boolean targets additionally go through the shared alias table so
    "yes"/"no"-style strings coerce the same way on both engines in best-effort
    mode, including Polars Categorical and Enum labels. Text-to-datetime casts
    share the pandas mixed-format parser and retain the resolved engine dtype.
    Columns missing from the frame are skipped. New categorical artifacts
    freeze the training vocabulary; unseen values become missing. Older
    artifacts without a vocabulary retain their original casting behavior.
    """

    @staticmethod
    def validate_inference_state(raw: dict) -> dict:
        """Inspect local dtype/vocabulary state without converting or relearning it."""
        fields = {"type", "type_map", "coerce_on_error"}
        # Optional fields: older artifacts predate them and keep their behavior.
        fields |= {"categories", "integral_floats_as_int"} & set(
            raw if isinstance(raw, dict) else ()
        )
        local_state_fields(raw, "casting", fields)
        _validate_casting_types(raw["type_map"])
        local_boolean(raw["coerce_on_error"], "coerce_on_error")
        if "integral_floats_as_int" in raw:
            local_boolean(raw["integral_floats_as_int"], "integral_floats_as_int")
        _validate_casting_categories(raw)
        return raw

    @staticmethod
    def inference_capability(state: dict, *, engine: str) -> ExecutionCapability | None:
        """Describe local casts, including request-derived legacy pandas vocabularies."""
        if engine not in ("pandas", "polars"):
            return None
        CastingApplier.validate_inference_state(state)
        context = _pandas_cast_context(state) if engine == "pandas" else "row"
        return ExecutionCapability(engine, "apply", "local", "preserve", context)

    @apply_method
    def apply(self, X: Any, _y: Any, params: dict[str, Any]) -> Any:  # pylint: disable=arguments-differ
        """Dispatch through the mapping-keyed dual-engine path; ``y`` passes through."""
        return apply_dual_engine(
            X, params, {"polars": _casting_apply_polars, "pandas": _casting_apply_pandas}
        )

@NodeRegistry.register("Casting", CastingApplier)
@node_meta(
    id="Casting",
    name="Type Casting",
    category="Data Operations",
    description="Cast columns to specific data types.",
    params={"column_types": {}, "coerce_on_error": True},
    learns_from_data=True,
)
class CastingCalculator(BaseCalculator):
    """Normalise either accepted config shape into one canonical ``type_map``.

    Ordinary casts resolve dtype aliases through ``TYPE_ALIASES``. Category
    casts additionally learn the observed training vocabulary so applying a
    held-out batch cannot change the accepted categories or pandas codes.
    """

    def infer_output_schema(
        self, input_schema: SkyulfSchema, config: dict[str, Any]
    ) -> SkyulfSchema:
        """Rewrite the dtype labels of ``input_schema`` for every column the config targets.

        The shared ``target_type`` overrides per-column choices, matching fit.
        Columns absent from the input schema are skipped, as they are at runtime.
        """
        # Casting preserves the column set but rewrites dtype labels.
        column_types = dict(config.get("column_types", {}) or {})
        target_type = config.get("target_type")
        columns = config.get("columns", []) or []
        if target_type and columns:
            for col in columns:
                column_types[col] = target_type
        new_schema = input_schema
        for col, dtype in column_types.items():
            resolved = _resolve_casting_dtype(dtype)
            new_schema = new_schema.with_dtype(col, str(resolved))
        return new_schema

    @fit_method
    def fit(self, X: Any, _y: Any, config: dict[str, Any]) -> CastingArtifact:  # pylint: disable=arguments-differ
        """Merge the two config shapes into a ``type_map`` of columns present on *X*.

        Absent columns are dropped rather than erroring, so a config written
        against a wider frame degrades to a partial cast; when both shapes name
        the same column, ``target_type`` overwrites the ``column_types`` entry.
        """
        # Config supports two shapes:
        #   {'columns': ['col1'], 'target_type': 'float'}
        #   {'column_types': {'col1': 'float', 'col2': 'int'}}
        target_type = config.get("target_type")
        columns = config.get("columns", [])
        column_types = config.get("column_types", {})

        final_map: dict[str, Any] = {}
        for col, dtype in column_types.items():
            if col in X.columns:
                final_map[col] = _resolve_casting_dtype(dtype)

        if target_type and columns:
            resolved_type = _resolve_casting_dtype(target_type)
            for col in columns:
                if col in X.columns:
                    final_map[col] = resolved_type

        categories = self._fit_categories(X, final_map)

        return {
            "type": "casting",
            "type_map": final_map,
            "coerce_on_error": config.get("coerce_on_error", True),
            "categories": categories,
            # New fits write 81.0 as "81" when casting to string; saved models replay as fitted.
            "integral_floats_as_int": True,
        }

    @staticmethod
    def _fit_categories(X: Any, type_map: dict[str, Any]) -> dict[str, list[Any]]:
        """Learn observed training categories without retaining unused vocabulary entries."""
        category_columns = [col for col, dtype in type_map.items() if dtype == "category"]
        categories: dict[str, list[Any]] = {}
        if category_columns:
            category_frame = select_then_to_pandas(X, category_columns)
            for col in category_columns:
                # Rebuild from observed values, excluding unused categories
                # that a sliced frame may have inherited from held-out rows.
                categories[col] = pd.Categorical(category_frame[col].to_numpy()).categories.tolist()

        return categories