# src/svy/wrangling/base.py
"""
Wrangling module for survey data transformation and cleaning.

Provides methods for:
- Column renaming and cleaning
- Value recoding and categorization
- Top/bottom coding (clamping)
- Row filtering
- Column creation and transformation (mutate)
- Label management
"""

from __future__ import annotations

import logging
import math

from typing import TYPE_CHECKING, Any, Callable, Literal, Mapping, Sequence, cast

import numpy as np
import polars as pl

from svy.core.constants import (
    SVY_HIT,
    SVY_PROB,
    SVY_ROW_INDEX,
    SVY_WEIGHT,
)
from svy.core.design import Design
from svy.core.enumerations import (
    CaseStyle,
    LetterCase,
    MeasurementType,
)
from svy.core.expr import Expr, to_polars_expr
from svy.core.types import MutateValue, WhereArg
from svy.engine.wrangling.cleaning import (
    _bottom_and_top_code,
    _bottom_code,
    _categorize,
    _clean_names,
    _recode,
    _rename,
    _top_code,
)
from svy.errors import DimensionError, LabelError, MethodError, SvyError
from svy.utils.helpers import (
    _normalize_columns_arg,
)


log = logging.getLogger(__name__)

INTEGER_DTYPES = {
    pl.Int8,
    pl.Int16,
    pl.Int32,
    pl.Int64,
    pl.UInt8,
    pl.UInt16,
    pl.UInt32,
    pl.UInt64,
}
FLOAT_DTYPES = {pl.Float32, pl.Float64}

# Scalar types that can be broadcast to all rows via pl.lit()
SCALAR_TYPES = (bool, int, float, str)

ExprLike = Expr | str | Callable[[Mapping[str, Expr]], Expr]

if TYPE_CHECKING:
    from svy.core.sample import Sample


# -------------------------------------------------------------------
# Helper utilities
# -------------------------------------------------------------------


def _check_nan_keys(mapping: dict, *, where: str, var: str | None = None) -> None:
    """Raise LabelError if any dict key is NaN."""
    for k in mapping.keys():
        try:
            if isinstance(k, float) and math.isnan(k):
                where_tag = f"{where} ({var})" if var else where
                raise LabelError.nan_key_forbidden(where=where_tag)
        except TypeError:
            pass


def _map_name_in_design(design_field: str | None, renames: dict[str, str]) -> str | None:
    """Maps a single string design field to a new name, or None if not found/string."""
    return (
        renames.get(design_field, design_field) if isinstance(design_field, str) else design_field
    )


def _map_tuple_in_design(
    design_field: str | tuple[str, ...] | None, renames: dict[str, str]
) -> str | tuple[str, ...] | None:
    """Maps a string or tuple design field to new name(s)."""
    if design_field is None or isinstance(design_field, str):
        return _map_name_in_design(design_field, renames)
    return tuple(renames.get(s, s) for s in design_field)


def _design_with_renamed_columns(design: Design, renames: dict[str, str]) -> Design:
    """
    Returns a new Design with all column references updated by `renames`.
    """
    if not renames:
        return design

    local_design = design

    # Map simple fields
    new_row_index = _map_name_in_design(local_design.row_index, renames)
    new_wgt = _map_name_in_design(local_design.wgt, renames)
    new_prob = _map_name_in_design(local_design.prob, renames)
    new_hit = _map_name_in_design(local_design.hit, renames)
    new_mos = _map_name_in_design(local_design.mos, renames)
    new_pop_size = _map_name_in_design(local_design.pop_size, renames)

    # Map Tuple-or-str fields
    new_stratum = _map_tuple_in_design(local_design.stratum, renames)
    new_psu = _map_tuple_in_design(local_design.psu, renames)
    new_ssu = _map_tuple_in_design(local_design.ssu, renames)

    # Replicate weights
    new_rep = local_design.rep_wgts
    if local_design.rep_wgts is not None:
        mapped_wgts = tuple(renames.get(s, s) for s in local_design.rep_wgts.wgts)  # type: ignore[union-attr]
        if mapped_wgts != local_design.rep_wgts.wgts:  # type: ignore[union-attr]
            new_rep = local_design.rep_wgts.clone(wgts=mapped_wgts, n_reps=len(mapped_wgts))  # type: ignore[union-attr]

    return local_design.update(
        row_index=new_row_index,
        stratum=new_stratum,
        wgt=new_wgt,
        prob=new_prob,
        hit=new_hit,
        mos=new_mos,
        psu=new_psu,
        ssu=new_ssu,
        pop_size=new_pop_size,
        rep_wgts=new_rep,
    )


def _as_output_obj(obj: object, out_name: str, n_rows: int) -> object:
    """
    Convert various input types to Polars expressions or Series for mutate().

    Handles:
    - None: creates null column
    - bool, int, float, str: broadcast scalar to all rows via pl.lit()
    - pl.Expr: Polars expression
    - svy.Expr: svy expression (converted to Polars)
    - pl.Series: renamed if length matches
    - np.ndarray, list, tuple: converted to Series if length matches

    Parameters
    ----------
    obj : object
        The value to convert.
    out_name : str
        The output column name.
    n_rows : int
        Expected number of rows (for length validation).

    Returns
    -------
    pl.Expr or pl.Series
        Object suitable for use with DataFrame.with_columns().

    Raises
    ------
    ValueError
        If array/Series length doesn't match n_rows.
    TypeError
        If obj type is not supported.
    """
    # Handle None -> null column
    if obj is None:
        return pl.lit(None).alias(out_name)

    # Handle Python scalar literals (broadcast to all rows)
    # Use type() to avoid matching numpy scalars or subclasses
    if type(obj) in SCALAR_TYPES:
        return pl.lit(obj).alias(out_name)

    # Handle Polars expression
    if isinstance(obj, pl.Expr):
        return obj.alias(out_name)

    # Handle svy expression
    if isinstance(obj, Expr):
        return to_polars_expr(obj).alias(out_name)

    # Handle Polars Series
    if isinstance(obj, pl.Series):
        if obj.len() != n_rows:
            raise ValueError(
                f"Series length {obj.len()} != n_rows {n_rows} for column '{out_name}'"
            )
        return obj.rename(out_name)

    # Handle numpy array
    try:
        is_np = isinstance(obj, np.ndarray)
    except Exception:
        is_np = False

    if is_np or isinstance(obj, (list, tuple)):
        vals = obj.tolist() if is_np else list(obj)  # type: ignore[arg-type]
        if len(vals) != n_rows:
            raise ValueError(
                f"Sequence length {len(vals)} != n_rows {n_rows} for column '{out_name}'"
            )
        return pl.Series(out_name, vals)

    # Fallback: try to convert to Polars expression
    try:
        return to_polars_expr(obj).alias(out_name)
    except Exception as ex:
        raise TypeError(
            f"Unsupported type for column '{out_name}': {type(obj).__name__}. "
            f"Expected: int, float, str, bool, None, Expr, pl.Expr, pl.Series, "
            f"np.ndarray, list, tuple, or callable."
        ) from ex


def _update_metadata_keys(sample: "Sample", renames: dict[str, str]) -> None:
    """Update metadata store keys after column rename."""
    if not renames:
        return

    meta = sample._metadata
    for old_name, new_name in renames.items():
        var_meta = meta.get(old_name)
        if var_meta is not None:
            meta.remove(old_name)
            meta.set(new_name, var_meta.clone(name=new_name))


# ---------------------------
# Wrangling facade on Sample
# ---------------------------


class Wrangling:
    """
    Data wrangling operations for survey samples.

    Accessed via `sample.wrangling` attribute. Provides methods for
    cleaning, transforming, and preparing survey data while maintaining
    design metadata consistency.

    All methods return the modified Sample, enabling method chaining.

    Examples
    --------
    >>> sample.wrangling.clean_names().wrangling.mutate({"x": 1})
    """

    def __init__(self, sample: Sample) -> None:
        self._sample = sample

    # ----- internal helpers -----

    def _internal_columns(self) -> set[str]:
        """
        Return set of internal svy columns that should be auto-preserved transparently.

        These columns are managed by svy internally and users should not need to
        think about them when using keep_columns/select or remove_columns/drop.
        """
        internal: set[str] = set()
        cols = set(self._sample._data.columns)

        # Always preserve svy_row_index if present
        if SVY_ROW_INDEX in cols:
            internal.add(SVY_ROW_INDEX)

        # Preserve any *_svy_internal_* columns (e.g., stratum_svy_internal_cols_concatenated)
        for c in cols:
            if "_svy_internal_" in c:
                internal.add(c)

        return internal

    def _required_columns(self) -> set[str]:
        """
        Return set of user-specified design column names.

        This excludes internal svy columns (like svy_row_index) which are
        handled transparently and auto-preserved.
        """
        req: set[str] = set()
        design: Design | None = getattr(self._sample, "_design", None)
        internal_design = getattr(self._sample, "_internal_design", {}) or {}

        def add(x: str | tuple[str, ...] | None):
            if x is None:
                return
            if isinstance(x, str):
                req.add(x)
            elif isinstance(x, tuple):
                req.update(x)

        if design is not None:
            # Note: row_index is NOT included here - it's an internal column
            add(design.stratum)
            add(design.psu)
            add(design.ssu)
            add(design.wgt)
            add(design.prob)
            add(design.hit)
            add(design.mos)
            add(design.pop_size)
            if design.rep_wgts and design.rep_wgts.wgts:
                req.update(design.rep_wgts.wgts)

        for k in ("stratum", "psu", "ssu"):
            cname = internal_design.get(k)
            if isinstance(cname, str):
                req.add(cname)

        cols = set(self._sample._data.columns)
        # Filter out internal columns from required set
        internal = self._internal_columns()
        return {c for c in req if c in cols and c not in internal}

    def _auto_clean_design(self) -> None:
        """Remove references to columns that no longer exist in the data."""
        current_design: Design | None = getattr(self._sample, "_design", None)
        if current_design is None:
            return

        cols = set(self._sample._data.columns)

        def keep_name(x: str | None) -> str | None:
            return x if (x is None or x in cols) else None

        def keep_tuple(x: str | tuple[str, ...] | None) -> str | tuple[str, ...] | None:
            if x is None:
                return None
            if isinstance(x, str):
                return x if x in cols else None
            kept = tuple(c for c in x if c in cols)
            return kept or None

        updated_design = current_design.update(
            row_index=keep_name(current_design.row_index),
            stratum=keep_tuple(current_design.stratum),
            psu=keep_tuple(current_design.psu),
            ssu=keep_tuple(current_design.ssu),
            wgt=keep_name(current_design.wgt),
            prob=keep_name(current_design.prob),
            hit=keep_name(current_design.hit),
            mos=keep_name(current_design.mos),
        )

        if current_design.rep_wgts:
            new_wgts = tuple(c for c in current_design.rep_wgts.wgts if c in cols)
            new_rep = (
                current_design.rep_wgts.clone(wgts=new_wgts, n_reps=len(new_wgts))
                if new_wgts
                else None
            )
            updated_design = updated_design.update(rep_wgts=new_rep)

        internal_design = dict(getattr(self._sample, "_internal_design", {}) or {})
        for k in ("stratum", "psu", "ssu"):
            cname = internal_design.get(k)
            if isinstance(cname, str) and cname not in cols:
                internal_design[k] = None

        self._sample._design = updated_design
        self._sample._internal_design = internal_design

    # ----- public API: Column Operations -----

    def clean_names(
        self,
        minimal: bool = False,
        remove: str | None = None,
        case_style: CaseStyle = CaseStyle.SNAKE,
        letter_case: LetterCase = LetterCase.LOWER,
    ) -> "Sample":
        """
        Standardize column names for easier downstream work.

        Parameters
        ----------
        minimal : bool, default False
            If True, only remove leading/trailing whitespace.
        remove : str, optional
            Regex pattern of characters to remove from column names.
        case_style : CaseStyle, default CaseStyle.SNAKE
            Target case style (SNAKE, CAMEL, PASCAL).
        letter_case : LetterCase, default LetterCase.LOWER
            Target letter case (LOWER, UPPER, TITLE).

        Returns
        -------
        Sample
            Sample with cleaned column names.

        Examples
        --------
        >>> sample.wrangling.clean_names(case_style=CaseStyle.SNAKE)
        >>> sample.wrangling.clean_names(remove=r"[^a-zA-Z0-9]")
        """
        _raw_data = self._sample._data
        _df_data: pl.DataFrame = (
            cast(pl.DataFrame, _raw_data)
            if not isinstance(_raw_data, pl.LazyFrame)
            else cast(pl.DataFrame, _raw_data.collect())
        )
        cleaned_data, renames = _clean_names(
            data=_df_data,
            minimal=minimal,
            remove=remove,
            case_style=case_style,
            letter_case=letter_case,
        )
        self._sample._data = cleaned_data
        if renames:
            _update_metadata_keys(self._sample, renames)
            if getattr(self._sample, "_design", None) is not None:
                self._sample._design = _design_with_renamed_columns(self._sample._design, renames)
        return self._sample

    def rename_columns(self, renames: dict[str, str]) -> "Sample":
        """
        Rename columns directly.

        Parameters
        ----------
        renames : dict[str, str]
            Mapping from old names to new names.

        Returns
        -------
        Sample
            Sample with renamed columns.

        Raises
        ------
        MethodError
            If trying to rename reserved svy columns or source name not found.

        Examples
        --------
        >>> sample.wrangling.rename_columns({"old_name": "new_name"})
        """
        if not renames:
            return self._sample

        forbidden = {SVY_ROW_INDEX, SVY_WEIGHT, SVY_PROB, SVY_HIT}
        if any(k in forbidden or v in forbidden for k, v in renames.items()):
            raise MethodError(
                title="Reserved column names cannot be renamed",
                detail=f"Forbidden: {sorted(forbidden)}",
                code="RENAME_FORBIDDEN",
                where="wrangling.rename_columns",
            )

        try:
            _raw2 = self._sample._data
            _df2: pl.DataFrame = (
                cast(pl.DataFrame, _raw2)
                if not isinstance(_raw2, pl.LazyFrame)
                else cast(pl.DataFrame, _raw2.collect())
            )
            self._sample._data = _rename(_df2, renames=renames)
        except (ValueError, KeyError):
            raise
        except Exception as ex:
            raise MethodError(
                title="Rename failed",
                detail=str(ex),
                code="RENAME_FAILED",
                where="wrangling.rename_columns",
            ) from ex

        _update_metadata_keys(self._sample, renames)

        if getattr(self._sample, "_design", None) is not None:
            self._sample._design = _design_with_renamed_columns(self._sample._design, renames)

        return self._sample

    def remove_columns(self, columns: str | Sequence[str], *, force: bool = False) -> "Sample":
        """
        Remove columns from the sample.

        Internal svy columns (like svy_row_index) cannot be removed and are
        silently preserved.

        Parameters
        ----------
        columns : str or Sequence[str]
            Column name(s) to remove.
        force : bool, default False
            If True, allow removing design-referenced columns.

        Returns
        -------
        Sample
            Sample with columns removed.

        Raises
        ------
        MethodError
            If trying to remove design columns without force=True.

        Examples
        --------
        >>> sample.wrangling.remove_columns(["temp_col", "unused"])
        """
        cols = _normalize_columns_arg(data=self._sample._data, columns=columns)

        # Filter out internal columns - they cannot be removed
        internal = self._internal_columns()
        cols = [c for c in cols if c not in internal]

        if not cols:
            return self._sample

        protected = self._required_columns()
        blocked = [c for c in cols if c in protected]
        if blocked and not force:
            raise MethodError(
                title="Cannot remove design-referenced columns",
                detail=", ".join(blocked),
                code="DROP_PROTECTED_COLUMNS",
                where="wrangling.remove_columns",
                hint="Pass force=True to drop them and automatically clean the design.",
            )
        self._sample._data = self._sample._data.drop(cols)
        if blocked:
            self._auto_clean_design()
        return self._sample

    def drop(self, columns: str | Sequence[str], *, force: bool = False) -> "Sample":
        """Alias for remove_columns()."""
        return self.remove_columns(columns, force=force)

    def keep_columns(self, columns: str | Sequence[str], *, force: bool = False) -> "Sample":
        """
        Keep only specified columns, removing all others.

        Internal svy columns (like svy_row_index) are automatically preserved
        and do not need to be included in the columns list.

        Parameters
        ----------
        columns : str or Sequence[str]
            Column name(s) to keep.
        force : bool, default False
            If True, allow dropping design-referenced columns.

        Returns
        -------
        Sample
            Sample with only specified columns (plus internal svy columns).

        Raises
        ------
        MethodError
            If keeping columns would remove design columns without force=True.

        Examples
        --------
        >>> sample.wrangling.keep_columns(["id", "weight", "outcome"])
        """
        cols = _normalize_columns_arg(data=self._sample._data, columns=columns)

        # Auto-include internal columns (transparent to user)
        internal = self._internal_columns()
        cols_with_internal = list(cols) + [c for c in internal if c not in cols]

        protected = self._required_columns()
        to_drop = [c for c in protected if c not in cols]
        if to_drop and not force:
            raise MethodError(
                title="Keeping these columns would remove design-referenced columns",
                detail=", ".join(to_drop),
                code="KEEP_DROPS_PROTECTED",
                where="wrangling.keep_columns",
                hint="Pass force=True to proceed; the design will be cleaned automatically.",
            )
        self._sample._data = self._sample._data.select(cols_with_internal)
        if to_drop:
            self._auto_clean_design()
        return self._sample

    def select(self, columns: str | Sequence[str], *, force: bool = False) -> "Sample":
        """Alias for keep_columns()."""
        return self.keep_columns(columns, force=force)

    # ----- public API: Value Transformation -----

    def top_code(
        self,
        top_codes: Mapping[str, float],
        *,
        replace: bool = False,
        into: str | Mapping[str, str] | None = None,
    ) -> "Sample":
        """
        Cap values at upper bounds (top coding).

        Parameters
        ----------
        top_codes : Mapping[str, float]
            Mapping from column names to maximum values.
        replace : bool, default False
            If True, replace original column. Otherwise create new column.
        into : str or Mapping[str, str], optional
            Custom output column name(s).

        Returns
        -------
        Sample
            Sample with capped values.

        Examples
        --------
        >>> sample.wrangling.top_code({"income": 200000})
        """
        _raw3 = self._sample._data
        _df3: pl.DataFrame = (
            cast(pl.DataFrame, _raw3)
            if not isinstance(_raw3, pl.LazyFrame)
            else cast(pl.DataFrame, _raw3.collect())
        )
        self._sample._data = _top_code(_df3, top_codes=top_codes, replace=replace, into=into)
        return self._sample

    def bottom_code(
        self,
        bottom_codes: Mapping[str, float],
        *,
        replace: bool = False,
        into: str | Mapping[str, str] | None = None,
    ) -> "Sample":
        """
        Cap values at lower bounds (bottom coding).

        Parameters
        ----------
        bottom_codes : Mapping[str, float]
            Mapping from column names to minimum values.
        replace : bool, default False
            If True, replace original column. Otherwise create new column.
        into : str or Mapping[str, str], optional
            Custom output column name(s).

        Returns
        -------
        Sample
            Sample with capped values.

        Examples
        --------
        >>> sample.wrangling.bottom_code({"age": 0})
        """
        _raw4 = self._sample._data
        _df4: pl.DataFrame = (
            cast(pl.DataFrame, _raw4)
            if not isinstance(_raw4, pl.LazyFrame)
            else cast(pl.DataFrame, _raw4.collect())
        )
        self._sample._data = _bottom_code(
            _df4, bottom_codes=bottom_codes, replace=replace, into=into
        )
        return self._sample

    def bottom_and_top_code(
        self,
        bottom_and_top_codes: Mapping[str, tuple[float, float] | list[float]],
        *,
        replace: bool = False,
        into: str | Mapping[str, str] | None = None,
    ) -> "Sample":
        """
        Cap values at both lower and upper bounds.

        Parameters
        ----------
        bottom_and_top_codes : Mapping[str, tuple[float, float]]
            Mapping from column names to (min, max) bounds.
        replace : bool, default False
            If True, replace original column. Otherwise create new column.
        into : str or Mapping[str, str], optional
            Custom output column name(s).

        Returns
        -------
        Sample
            Sample with clamped values.

        Examples
        --------
        >>> sample.wrangling.bottom_and_top_code({"income": (0, 200000)})
        """
        _raw5 = self._sample._data
        _df5: pl.DataFrame = (
            cast(pl.DataFrame, _raw5)
            if not isinstance(_raw5, pl.LazyFrame)
            else cast(pl.DataFrame, _raw5.collect())
        )
        self._sample._data = _bottom_and_top_code(
            _df5,
            bottom_and_top_codes=bottom_and_top_codes,
            replace=replace,
            into=into,
        )
        return self._sample

    def recode(
        self,
        cols: str | list[str],
        recodes: Mapping[Any, Sequence[Any]],
        *,
        replace: bool = False,
        into: Mapping[str, str] | str | None = None,
    ) -> "Sample":
        """
        Map old values to new labels.

        Values not in the mapping pass through unchanged.

        Parameters
        ----------
        cols : str or list[str]
            Column(s) to recode.
        recodes : Mapping[Any, Sequence[Any]]
            Mapping from new value to list of old values.
            E.g., {"Food": ["apple", "carrot"]} maps both to "Food".
        replace : bool, default False
            If True, replace original column.
        into : str or Mapping[str, str], optional
            Custom output column name(s).

        Returns
        -------
        Sample
            Sample with recoded values.

        Examples
        --------
        >>> sample.wrangling.recode("item", {"Food": ["apple", "carrot"]})
        """
        _raw6 = self._sample._data
        _df6: pl.DataFrame = (
            cast(pl.DataFrame, _raw6)
            if not isinstance(_raw6, pl.LazyFrame)
            else cast(pl.DataFrame, _raw6.collect())
        )
        self._sample._data = _recode(_df6, cols=cols, recodes=recodes, replace=replace, into=into)
        return self._sample

    def categorize(
        self,
        col: str,
        bins: list[float],
        labels: list[str] | None = None,
        *,
        right: bool = True,
        replace: bool = False,
        into: str | None = None,
    ) -> "Sample":
        """
        Bin continuous values into labeled categories.

        Out-of-range values become null.

        Parameters
        ----------
        col : str
            Column to categorize.
        bins : list[float]
            Bin edges. Creates len(bins)-1 categories.
        labels : list[str], optional
            Labels for each bin. Must have len(bins)-1 elements.
        right : bool, default True
            If True, bins are right-closed (a, b]. If False, left-closed [a, b).
        replace : bool, default False
            If True, replace original column.
        into : str, optional
            Custom output column name.

        Returns
        -------
        Sample
            Sample with categorized values.

        Examples
        --------
        >>> sample.wrangling.categorize(
        ...     "age",
        ...     bins=[0, 18, 65, 100],
        ...     labels=["child", "adult", "senior"]
        ... )
        """
        _raw7 = self._sample._data
        _df7: pl.DataFrame = (
            cast(pl.DataFrame, _raw7)
            if not isinstance(_raw7, pl.LazyFrame)
            else cast(pl.DataFrame, _raw7.collect())
        )
        self._sample._data = _categorize(
            data=_df7,
            varname=col,
            bins=bins,
            labels=labels,
            right=right,
            replace=replace,
            into=into,
        )
        return self._sample

    # ----- public API: Data Type Operations -----

    def cast(
        self,
        cols: str | Sequence[str] | Mapping[str, pl.DataType],
        dtype: pl.DataType | None = None,
        *,
        strict: bool = True,
    ) -> "Sample":
        """
        Cast columns to specified data type(s).

        Parameters
        ----------
        cols : str, Sequence[str], or Mapping[str, DataType]
            Column(s) to cast. If a mapping, values are target dtypes.
        dtype : DataType, optional
            Target data type (required if cols is not a mapping).
        strict : bool, default True
            If True, raise on invalid casts. If False, return null.

        Returns
        -------
        Sample
            Sample with cast columns.

        Examples
        --------
        >>> sample.wrangling.cast("age", pl.Float64)
        >>> sample.wrangling.cast({"age": pl.Float64, "id": pl.Utf8})
        """

        def _cast_expr(col_name: str, target: pl.DataType, strict: bool) -> pl.Expr:
            """Build cast expression, routing through Utf8 for Categorical/Enum."""
            base = pl.col(col_name)
            if target in (pl.Categorical, pl.Enum) or (
                isinstance(target, pl.Categorical) or isinstance(target, pl.Enum)
            ):
                return base.cast(pl.Utf8).cast(target, strict=strict).alias(col_name)
            return base.cast(target, strict=strict).alias(col_name)

        if isinstance(cols, Mapping):
            exprs = [_cast_expr(c, dt, strict) for c, dt in cols.items()]  # type: ignore[arg-type]
        else:
            if dtype is None:
                raise MethodError(
                    title="dtype required when cols is not a mapping",
                    detail="Provide dtype parameter or pass a dict mapping columns to types",
                    code="CAST_DTYPE_REQUIRED",
                    where="wrangling.cast",
                )
            col_list = [cols] if isinstance(cols, str) else list(cols)
            exprs = [_cast_expr(c, dtype, strict) for c in col_list]

        self._sample._data = self._sample._data.with_columns(exprs)
        return self._sample

    def fill_null(
        self,
        cols: str | Sequence[str],
        value: Any = None,
        *,
        strategy: Literal["forward", "backward", "mean", "min", "max", "zero", "one"]
        | None = None,
    ) -> "Sample":
        """
        Fill null values in specified columns.

        Parameters
        ----------
        cols : str or Sequence[str]
            Column(s) to fill.
        value : Any, optional
            Value to fill nulls with (ignored if strategy is set).
        strategy : str, optional
            Fill strategy: "forward", "backward", "mean",
            "min", "max", "zero", "one".

        Returns
        -------
        Sample
            Sample with filled nulls.

        Examples
        --------
        >>> sample.wrangling.fill_null("income", value=0)
        >>> sample.wrangling.fill_null("income", strategy="mean")
        """
        col_list = [cols] if isinstance(cols, str) else list(cols)

        if strategy is not None:
            exprs = [pl.col(c).fill_null(strategy=strategy).alias(c) for c in col_list]
        else:
            exprs = [pl.col(c).fill_null(value).alias(c) for c in col_list]

        self._sample._data = self._sample._data.with_columns(exprs)
        return self._sample

    # ----- public API: Row Operations -----

    def filter_records(
        self,
        where: WhereArg | None = None,
        *,
        negate: bool = False,
        check_singletons: bool = False,
        on_singletons: Literal["ignore", "warn", "error"] = "ignore",
    ) -> "Sample":
        """
        Filter rows based on conditions.

        Parameters
        ----------
        where : WhereArg, optional
            Filter condition. Can be:
            - dict: {col: value} or {col: [values]} for equality/membership
            - list of Expr: combined with AND
            - single Expr
        negate : bool, default False
            If True, keep rows that do NOT match the condition.
        check_singletons : bool, default False
            If True, check for singleton PSUs after filtering.
        on_singletons : {"ignore", "warn", "error"}, default "ignore"
            Action if singletons are detected.

        Returns
        -------
        Sample
            Filtered sample.

        Examples
        --------
        >>> sample.wrangling.filter_records({"region": "North"})
        >>> sample.wrangling.filter_records([svy.col("age") > 18])
        """
        if where is None:
            return self._sample

        try:
            pred = self._compile_where_to_pl_expr(where)
            if pred is None:
                return self._sample
            if negate:
                pred = ~pred

            self._sample._data = self._sample._data.filter(pred)

            if check_singletons:
                self._sample._check_for_singletons()
                if getattr(self._sample, "_singletons", None):
                    if on_singletons == "warn":
                        self._sample.warn(
                            code="SINGLETONS_DETECTED",
                            title="Singleton PSUs/strata detected after filtering",
                            detail=f"Found {len(self._sample._singletons or [])} singleton group(s).",
                            where="wrangling.filter_records",
                        )
                    elif on_singletons == "error":
                        raise MethodError(
                            title="Singletons detected after filtering",
                            detail=f"Found {len(self._sample._singletons or [])} singleton group(s).",
                            code="SINGLETONS_AFTER_FILTER",
                            where="wrangling.filter_records",
                            hint="Collapse strata, adjust PSUs, or handle via the singleton utilities.",
                        )

            return self._sample

        except SvyError:
            raise
        except Exception as ex:
            raise MethodError(
                title="Filter failed",
                detail=str(ex),
                code="FILTER_FAILED",
                where="wrangling.filter_records",
            ) from ex

    def order_by(
        self,
        cols: str | Sequence[str],
        *,
        descending: bool | Sequence[bool] = False,
        nulls_last: bool = True,
    ) -> "Sample":
        self._sample._data = self._sample._data.sort(
            by=cols, descending=descending, nulls_last=nulls_last
        )
        return self._sample

    def distinct(
        self,
        cols: str | Sequence[str] | None = None,
        *,
        keep: Literal["first", "last", "any", "none"] = "first",
        maintain_order: bool = True,
    ) -> "Sample":
        """
        Remove duplicate rows.

        Parameters
        ----------
        cols : str, Sequence[str], or None
            Column(s) to consider for uniqueness. If None, all columns.
        keep : {"first", "last", "any", "none"}, default "first"
            Which duplicate to keep.
        maintain_order : bool, default True
            If True, preserve row order.

        Returns
        -------
        Sample
            Sample with duplicates removed.

        Examples
        --------
        >>> sample.wrangling.distinct("household_id")
        """
        subset = [cols] if isinstance(cols, str) else cols
        self._sample._data = self._sample._data.unique(
            subset=subset,
            keep=keep,
            maintain_order=maintain_order,
        )
        return self._sample

    def with_row_index(self, name: str = "row_index", offset: int = 0) -> "Sample":
        """
        Add a row index column.

        Parameters
        ----------
        name : str, default "row_index"
            Name for the index column.
        offset : int, default 0
            Starting value for the index. Cannot be negative.

        Returns
        -------
        Sample
            Sample with row index column added.

        Examples
        --------
        >>> sample.wrangling.with_row_index("id", offset=1)
        """
        self._sample._data = self._sample._data.with_row_index(name=name, offset=offset)
        return self._sample

    # ----- private: compile where -> pl.Expr -----

    def _compile_where_to_pl_expr(self, where: WhereArg) -> pl.Expr | None:
        """Convert WhereArg to a Polars filter expression."""
        if isinstance(where, Mapping):
            preds: list[pl.Expr] = []
            for k, v in where.items():
                if isinstance(v, (list, tuple, set)) and not isinstance(
                    v, (str, bytes, bytearray)
                ):
                    preds.append(pl.col(k).is_in(list(v)))
                else:
                    preds.append(pl.col(k) == v)  # type: ignore[arg-type]
            if not preds:
                return None
            acc = preds[0]
            for p in preds[1:]:
                acc = acc & p
            return acc

        if isinstance(where, (list, tuple)) and not isinstance(where, (str, bytes, bytearray)):
            if not where:
                return None
            compiled = [to_polars_expr(e) for e in where]
            acc = compiled[0]
            for p in compiled[1:]:
                acc = acc & p
            return acc

        return to_polars_expr(where)

    # ----- public API: Mutate -----

    def mutate(
        self,
        specs: Mapping[str, MutateValue],
        *,
        inplace: bool = False,
    ) -> "Sample":
        """
        Create or transform columns using expressions, scalars, or arrays.

        Parameters
        ----------
        specs : Mapping[str, MutateValue]
            Dictionary mapping output column names to values. Values can be:

            - **Scalars** (int, float, str, bool, None): broadcast to all rows
            - **svy.Expr** or **pl.Expr**: Polars expressions
            - **np.ndarray, list, tuple**: must match n_rows
            - **pl.Series**: must match n_rows
            - **Callable**: receives column namespace dict, returns expression

        inplace : bool, default False
            If True, modify the sample in place. If False, return a new sample.

        Returns
        -------
        Sample
            Modified sample with new/transformed columns.

        Notes
        -----
        - Columns created in the same call can reference each other
        - Circular dependencies are detected and raise an error
        - Use `svy.when().then().otherwise()` for conditional logic

        Examples
        --------
        >>> # Scalar values
        >>> sample.wrangling.mutate({"_one": 1, "pi": 3.14159})

        >>> # Expressions
        >>> sample.wrangling.mutate({
        ...     "income_k": svy.col("income") / 1000,
        ...     "log_income": svy.col("income").log(),
        ... })

        >>> # Conditional logic
        >>> sample.wrangling.mutate({
        ...     "age_group": svy.when(svy.col("age") < 18).then("child")
        ...                    .otherwise("adult")
        ... })

        >>> # Dependencies within same call
        >>> sample.wrangling.mutate({
        ...     "a": svy.col("x") * 2,
        ...     "b": svy.col("a") + 1,  # References 'a' created above
        ... })
        """
        _raw_ld = self._sample._data
        local_data: pl.DataFrame = (
            cast(pl.DataFrame, _raw_ld)
            if not isinstance(_raw_ld, pl.LazyFrame)
            else cast(pl.DataFrame, _raw_ld.collect())
        )
        n_rows = local_data.height

        env = {name: pl.col(name) for name in local_data.columns}

        compiled: dict[str, tuple[object, set[str]]] = {}

        def _root_names_safe(expr: pl.Expr) -> set[str]:
            try:
                return set(expr.meta.root_names())
            except Exception:
                return set()

        for out_name, spec in specs.items():
            try:
                raw = spec(env) if callable(spec) else spec  # type: ignore[operator]
                out_obj = _as_output_obj(raw, out_name, n_rows)
            except Exception as ex:
                raise MethodError(
                    title="Column transformation failed",
                    detail=str(ex),
                    code="MUTATE_COMPILE_FAILED",
                    where="wrangling.mutate",
                    hint="Check names, syntax, that array/series lengths match n_rows, and finalize when/then with .otherwise(...).",
                ) from ex

            deps: set[str] = _root_names_safe(out_obj) if isinstance(out_obj, pl.Expr) else set()
            compiled[out_name] = (out_obj, deps)

        # Batch by readiness (topological sort)
        current_cols = set(local_data.columns)
        pending = set(compiled.keys())
        max_iters = len(pending) + 8

        for _ in range(max_iters):
            ready = [name for name in pending if compiled[name][1] <= current_cols]
            if not ready:
                break

            batch_objs = [compiled[name][0] for name in ready]
            try:
                local_data = local_data.with_columns(batch_objs)
            except pl.exceptions.ColumnNotFoundError as ex:
                missing_by_out = {
                    name: sorted(compiled[name][1] - current_cols)
                    for name in ready
                    if compiled[name][1] - current_cols
                }
                raise DimensionError.missing_columns(
                    missing=sum(missing_by_out.values(), []),
                    param="columns",
                    where="wrangling.mutate",
                ) from ex
            except Exception as ex:
                raise MethodError(
                    title="Column transformation failed",
                    detail=str(ex),
                    code="MUTATE_WITH_COLUMNS_FAILED",
                    where="wrangling.mutate",
                ) from ex

            for name in ready:
                pending.remove(name)
                current_cols.add(name)

            if not pending:
                break

        if pending:
            unresolved = {name: sorted(compiled[name][1] - current_cols) for name in pending}
            missing_outside = sorted(
                {
                    col
                    for deps in unresolved.values()
                    for col in deps
                    if col not in pending and col not in current_cols
                }
            )

            if missing_outside:
                raise DimensionError.missing_columns(
                    where="wrangling.mutate",
                    param="specs",
                    missing=missing_outside,
                    available=local_data.columns,
                )

            raise MethodError(
                title="Dependency resolution failed",
                detail=f"Circular or mutually dependent outputs: {unresolved}",
                code="MUTATE_DEPENDENCY_ERROR",
                where="wrangling.mutate",
                hint="Split into multiple mutate() calls or break the circular reference.",
            )

        # Return result based on inplace flag
        if inplace:
            self._sample._data = local_data
            return self._sample
        else:
            return self._sample._replace_data(local_data)

    # ----- public API: Labels -----

    def apply_labels(
        self,
        cols: str | list[str],
        labels: str | list[str],
        categories: dict | None = None,
        *,
        strict: bool = True,
        overwrite: bool = True,
    ) -> "Sample":
        """
        Apply variable labels and optional value labels to columns.

        Parameters
        ----------
        cols : str or list[str]
            Column name(s) to label.
        labels : str or list[str]
            Label(s) for each column.
        categories : dict, optional
            Value labels. Can be:
            - Global: {value: label} applied to all cols
            - Per-variable: {col: {value: label}}
        strict : bool, default True
            If True, raise if column not found.
        overwrite : bool, default True
            If True, replace existing labels.

        Returns
        -------
        Sample
            Sample with labels applied.

        Examples
        --------
        >>> sample.wrangling.apply_labels("sex", "Sex of respondent", {1: "Male", 2: "Female"})
        """
        meta = self._sample._metadata

        vs = [cols] if isinstance(cols, str) else list(cols)
        ls = [labels] if isinstance(labels, str) else list(labels)

        if len(vs) != len(ls):
            raise MethodError(
                title="Label assignment mismatch",
                detail=f"Lengths differ: vars={len(vs)} labels={len(ls)}",
                code="LABELS_LENGTH_MISMATCH",
                where="wrangling.apply_labels",
                expected="len(vars) == len(labels)",
                got={"vars": len(vs), "labels": len(ls)},
                hint="Provide one label per variable name.",
            )

        local_data = self._sample._data
        col_names = set(local_data.columns)
        missing = [v for v in vs if v not in col_names]
        if missing:
            if strict:
                raise DimensionError.missing_columns(
                    where="wrangling.apply_labels",
                    param="vars",
                    missing=missing,
                    available=local_data.columns,
                )
            keep = [i for i, v in enumerate(vs) if v in col_names]
            vs = [vs[i] for i in keep]
            ls = [ls[i] for i in keep]
            if not vs:
                return self._sample

        per_var_categories: dict[str, dict] | None = None
        global_categories: dict | None = None

        if categories is not None:
            if not isinstance(categories, dict):
                raise MethodError.invalid_type(
                    where="wrangling.apply_labels",
                    param="categories",
                    got=categories,
                    expected="dict[Category, str] or dict[str, dict[Category, str]]",
                    hint="Use {'var': {0:'No',1:'Yes'}} or {0:'No',1:'Yes'}.",
                )
            if any(isinstance(v, dict) for v in categories.values()):
                per_var_categories = categories
            else:
                global_categories = categories

            if global_categories is not None:
                _check_nan_keys(global_categories, where="wrangling.apply_labels:global")
            else:
                for var_name, m in per_var_categories.items():  # type: ignore[union-attr]
                    if m is not None:
                        if not isinstance(m, dict):
                            raise MethodError.invalid_type(
                                where="wrangling.apply_labels",
                                param=f"categories[{var_name}]",
                                got=type(m).__name__,
                                expected="dict[Category, str] or None",
                            )
                        _check_nan_keys(m, where="wrangling.apply_labels", var=var_name)

        allowed_for_value_labels = {
            MeasurementType.NOMINAL,
            MeasurementType.ORDINAL,
            MeasurementType.BOOLEAN,
        }

        def _measurement_of(var: str):
            meas = None
            schema = getattr(self._sample, "_schema", None)
            if schema is not None:
                try:
                    meas = getattr(schema[var], "measurement", None)
                except Exception:
                    pass
            if meas is None:
                mget = getattr(self._sample, "measurement_of", None)
                if callable(mget):
                    try:
                        meas = mget(var)
                    except Exception:
                        pass
            return meas

        for v, lb in zip(vs, ls):
            # Check for existing labels when overwrite=False
            existing = meta.get(v)
            if not overwrite and existing is not None and existing.label is not None:
                raise MethodError.not_applicable(
                    where="wrangling.apply_labels",
                    method="apply_labels",
                    reason=f"variable '{v}' already has labels; set overwrite=True to replace",
                    param=v,
                    hint="Pass overwrite=True to replace existing label mapping.",
                )

            mapping_for_v = None
            if per_var_categories is not None:
                mapping_for_v = per_var_categories.get(v)
            elif global_categories is not None:
                mapping_for_v = global_categories

            if mapping_for_v:
                meas = _measurement_of(v)
                if meas is not None and meas not in allowed_for_value_labels:
                    raise MethodError.not_applicable(
                        where="wrangling.apply_labels",
                        method="value labels",
                        reason=f"variable '{v}' is measured as {getattr(meas, 'name', meas)}",
                        param=v,
                        hint="Value labels are supported for NOMINAL/ORDINAL/BOOLEAN only. Remove `categories` or convert the measurement type.",
                    )

                for key, val in mapping_for_v.items():
                    if not isinstance(val, str):
                        raise MethodError.invalid_type(
                            where="wrangling.apply_labels",
                            param=f"categories[{v}][{key}]",
                            got=type(val).__name__,
                            expected="str",
                            hint="Value labels must be strings.",
                        )

                # Set both variable label and value labels
                meta.set_label(v, lb)
                meta.set_value_labels(v, dict(mapping_for_v))
            else:
                # Set only variable label
                meta.set_label(v, lb)

        return self._sample
