# src/svy/weighting/base.py
from __future__ import annotations

import logging

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

import numpy as np
import polars as pl


try:
    from svy_rs._internal import (  # type: ignore[import-untyped]
        adjust_nr as rust_adjust_nr,
    )
    from svy_rs._internal import (  # type: ignore[import-untyped]
        calibrate as rust_calibrate,
    )
    from svy_rs._internal import (  # type: ignore[import-untyped]
        calibrate_by_domain as rust_calibrate_by_domain,
    )
    from svy_rs._internal import (  # type: ignore[import-untyped]
        calibrate_parallel as rust_calibrate_parallel,
    )
    from svy_rs._internal import (  # type: ignore[import-untyped]
        create_bootstrap_wgts as rust_create_bootstrap_wgts,
    )
    from svy_rs._internal import (  # type: ignore[import-untyped]
        create_brr_wgts as rust_create_brr_wgts,
    )
    from svy_rs._internal import (  # type: ignore[import-untyped]
        create_jk_wgts as rust_create_jk_wgts,
    )
    from svy_rs._internal import (  # type: ignore[import-untyped]
        create_sdr_wgts as rust_create_sdr_wgts,
    )
    from svy_rs._internal import (  # type: ignore[import-untyped]
        normalize as rust_normalize,
    )
    from svy_rs._internal import (  # type: ignore[import-untyped]
        poststratify as rust_poststratify,
    )
    from svy_rs._internal import (  # type: ignore[import-untyped]
        rake as rust_rake,
    )
except ImportError:  # pragma: no cover
    rust_adjust_nr = rust_calibrate = rust_calibrate_by_domain = None
    rust_calibrate_parallel = rust_create_bootstrap_wgts = None
    rust_create_brr_wgts = rust_create_jk_wgts = rust_create_sdr_wgts = None
    rust_normalize = rust_poststratify = rust_rake = None

from svy.core.constants import (
    SVY_ADJ_PREFIX,
    SVY_CALIB_PREFIX,
    SVY_NORM_PREFIX,
    SVY_PS_PREFIX,
    SVY_RAKED_PREFIX,
)
from svy.core.design import RepWeights
from svy.core.enumerations import EstimationMethod
from svy.core.terms import Cat, Cross, Feature
from svy.core.types import (
    Category,
    ControlsType,
    DomainScalarMap,
    NDArray,
    Number,
)
from svy.engine.weighting.adj_calibration import _calibrate
from svy.engine.weighting.adj_nonresponse import _adjust_nr
from svy.engine.weighting.adj_normalization import _normalize, _normalize_controls_like
from svy.engine.weighting.adj_poststratification import _poststratify
from svy.engine.weighting.adj_raking import _calculate_controls_from_factors, _rake
from svy.errors import DimensionError, MethodError
from svy.utils.checks import drop_missing
from svy.utils.random_state import RandomState, resolve_random_state


if TYPE_CHECKING:
    from svy.core.sample import Sample

log = logging.getLogger(__name__)


# ---------------------------
# Internal Helpers
# ---------------------------


def _to_int_array(
    data: pl.DataFrame,
    col_spec: str | tuple[str, ...] | None,
) -> np.ndarray | None:
    """
    Convert a single or multi-column design field to a 1-D int64 array
    suitable for passing to Rust functions.

    Parameters
    ----------
    data     : source DataFrame
    col_spec : str, tuple[str, ...], or None
               - None          → returns None
               - str           → encodes that single column
               - tuple[str,…]  → concatenates columns into composite labels,
                                  then integer-encodes

    Returns
    -------
    np.ndarray[int64] | None
    """
    if col_spec is None:
        return None
    cols = [col_spec] if isinstance(col_spec, str) else list(col_spec)
    if len(cols) == 1:
        vec = data[cols[0]].to_numpy()
    else:
        # Concatenate multi-column keys into a single composite string label
        vec = np.array(
            ["__".join(str(v) for v in row) for row in data.select(cols).iter_rows()],
            dtype=object,
        )
    if vec.dtype.kind in ("U", "S", "O"):
        _, codes = np.unique(vec, return_inverse=True)
        return codes.astype(np.int64)
    return vec.astype(np.int64)


def _to_float_array(
    data: pl.DataFrame,
    col: str | None,
    n: int,
) -> np.ndarray:
    """
    Return a float64 weight array from a column, or ones if col is None.
    Always returns float64 regardless of source dtype.
    """
    if col and col in data.columns:
        return data[col].to_numpy().astype(np.float64, copy=False)
    return np.ones(n, dtype=np.float64)


def _canonical_psu_table(df: pl.DataFrame, by_cols: list[str]) -> pl.DataFrame:
    return df.select(by_cols).unique(maintain_order=True).sort(by_cols)


def _name_rep_cols(prefix: str, n_reps: int) -> list[str]:
    return [f"{prefix}{i}" for i in range(1, n_reps + 1)]


def _ensure_columns(df: pl.DataFrame, cols: list[str], *, where: str) -> None:
    missing = [c for c in cols if c not in df.columns]
    if missing:
        raise MethodError.invalid_choice(
            where=where,
            param="columns",
            got=missing,
            allowed=df.columns,
            hint="Check spelling or pass existing column names.",
        )


def _unique_name(base: str, existing: set[str]) -> str:
    if base not in existing:
        return base
    k = 1
    while f"{base}_{k}" in existing:
        k += 1
    return f"{base}_{k}"


def _extract_cols_from_term(term: Feature) -> list[str]:
    if isinstance(term, str):
        return [term]
    if isinstance(term, Cat):
        return [term.name]
    if isinstance(term, Cross):
        return _extract_cols_from_term(term.left) + _extract_cols_from_term(term.right)
    raise TypeError(f"Unsupported term type for grouping: {type(term)}")


def _normalize_by_term(
    df: pl.DataFrame, by: Feature | None, *, param_name: str = "by", where: str
) -> tuple[list[str] | None, NDArray | None]:
    if by is None:
        return None, None

    if isinstance(by, (list, tuple)):
        raise TypeError(
            f"`{param_name}` must be a single Feature (str, Cat, Cross). "
            "Lists are not allowed; use Cross() for interactions."
        )

    adj_cols = _extract_cols_from_term(by)

    missing = [c for c in adj_cols if c not in df.columns]
    if missing:
        raise TypeError(f"`{param_name}` term refers to columns not present in data: {missing!r}")

    if len(adj_cols) == 1:
        adj_class_arr = df.get_column(adj_cols[0]).to_numpy()
    else:
        cols_np = [df.get_column(c).to_numpy() for c in adj_cols]
        adj_class_arr = np.fromiter(
            (tuple(row) for row in zip(*cols_np)),
            dtype=object,
            count=df.height,
        )

    return adj_cols, adj_class_arr


def _num_sort_key_token(tok: str) -> tuple[int, float | str]:
    try:
        return (0, float(tok))
    except Exception:
        return (1, tok)


def _num_sort_key_label(lbl):
    if isinstance(lbl, tuple):
        return tuple(_num_sort_key_token(t) for t in lbl)
    return _num_sort_key_token(lbl)


# ---------------------------
# Term Expansion & Target Matching (for Calibration)
# ---------------------------


def _expand_term(
    term: Feature, df: pl.DataFrame, where: str
) -> tuple[list[pl.Expr], list[Category]]:
    if isinstance(term, str):
        if term not in df.columns:
            raise KeyError(f"{where}: Continuous term '{term}' not found in data.")
        return [pl.col(term).cast(pl.Float64).fill_null(0.0)], [term]

    if isinstance(term, Cat):
        col_name = term.name
        if col_name not in df.columns:
            raise KeyError(f"{where}: Cat variable '{col_name}' not found.")

        levels = df.get_column(col_name).unique().sort().to_list()

        if term.ref is not None:
            if term.ref not in levels:
                raise ValueError(
                    f"{where}: Reference level '{term.ref}' not found in '{col_name}'."
                )
            levels = [lbl for lbl in levels if lbl != term.ref]

        exprs = []
        labels = []
        for i, lvl in enumerate(levels):
            if lvl is None:
                base_expr = pl.col(col_name).is_null()
            else:
                base_expr = pl.col(col_name) == lvl

            expr = base_expr.cast(pl.Float64).fill_null(0.0).alias(f"_tmp_cat_{col_name}_{i}")
            exprs.append(expr)
            labels.append(lvl)

        return exprs, labels

    if isinstance(term, Cross):
        left_exprs, left_labs = _expand_term(term.left, df, where)
        right_exprs, right_labs = _expand_term(term.right, df, where)

        out_exprs = []
        out_labs = []

        count = 0
        for le, ll in zip(left_exprs, left_labs):
            for re, rl in zip(right_exprs, right_labs):
                out_exprs.append((le * re).alias(f"_tmp_cross_{count}"))
                count += 1

                def _to_tuple(x):
                    return x if isinstance(x, tuple) else (x,)

                new_lab = _to_tuple(ll) + _to_tuple(rl)
                out_labs.append(new_lab)

        return out_exprs, out_labs

    raise TypeError(f"Unsupported term type: {type(term)}")


def _match_term_targets(
    labels: list[Category],
    target_spec: Number | dict[Category, Number] | Sequence[Number],
    term_name: str,
) -> list[float]:
    if isinstance(target_spec, (int, float, np.integer, np.floating)):
        if len(labels) != 1:
            raise ValueError(
                f"Scalar target provided for term '{term_name}' which expanded "
                f"to {len(labels)} columns. Use a dict mapping labels to values."
            )
        return [float(target_spec)]

    if isinstance(target_spec, dict):
        targets = []
        missing = [lbl for lbl in labels if lbl not in target_spec]
        if missing:
            raise ValueError(f"Missing targets for term '{term_name}': {missing}")

        for lbl in labels:
            val = target_spec[lbl]  # type: ignore[index]
            if not isinstance(val, (int, float, np.integer, np.floating)):
                raise TypeError(f"Target for '{term_name}':{lbl} must be numeric.")
            targets.append(float(val))
        return targets

    raise TypeError(
        f"Unsupported target specification type for '{term_name}': {type(target_spec)}"
    )


# ---------------------------
# Weighting Facade
# ---------------------------


class Weighting:
    def __init__(self, sample: Any) -> None:
        self._sample = sample

    def create_variance_strata(
        self,
        *,
        method: Literal["brr", "jk2"],
        order_by: str | Sequence[str] | None = None,
        shuffle: bool = False,
        into: str = "svy_var_stratum",
        rstate: int | None = None,
    ) -> Any:
        """
        Create variance estimation strata for BRR or JK2.

        Converts a design with multiple PSUs per stratum into one suitable
        for BRR (exactly 2 PSUs per variance stratum) or JK2 (2-3 PSUs).

        Args:
            method: Variance estimation method:
                - "brr": Creates pairs (2 PSUs each). Requires even PSU counts.
                - "jk2": Creates pairs, allows triplet for odd counts (2-3 PSUs each).
            order_by: Column(s) to sort PSUs by before pairing.
            shuffle: If True, randomly shuffle PSUs before pairing.
            into: Name for the new variance stratum column.
            rstate: Random seed for shuffle=True.

        Returns:
            Sample with new variance stratum column added and design updated.
        """
        where = "Sample.weighting.create_variance_strata"
        df = self._sample._data
        design = self._sample._design

        if design.psu is None:
            raise MethodError.not_applicable(
                where=where,
                method="create_variance_strata",
                reason="Design must have PSU defined.",
            )

        _method = method.lower()
        if _method not in ("brr", "jk2"):
            raise MethodError.invalid_choice(
                where=where,
                param="method",
                got=_method,
                allowed=["brr", "jk2"],
            )

        psu_col = design.psu
        orig_stratum_col = design.stratum

        if order_by is None:
            order_cols = []
        elif isinstance(order_by, str):
            order_cols = [order_by]
        else:
            order_cols = list(order_by)

        missing_cols = [col for col in order_cols if col not in df.columns]
        if missing_cols:
            raise MethodError.invalid_choice(
                where=where,
                param="order_by",
                got=missing_cols,
                allowed=df.columns,
                hint="Check column names.",
            )

        select_cols = [psu_col]
        if orig_stratum_col:
            select_cols.append(orig_stratum_col)
        select_cols.extend([c for c in order_cols if c not in select_cols])

        psu_df = df.select(select_cols).unique(maintain_order=True)

        sort_cols = []
        if orig_stratum_col:
            sort_cols.append(orig_stratum_col)
        sort_cols.extend(order_cols)
        if psu_col not in sort_cols:
            sort_cols.append(psu_col)

        psu_df = psu_df.sort(sort_cols)

        psu_list = psu_df[psu_col].to_list()
        n_psus = len(psu_list)

        if orig_stratum_col:
            orig_strata = psu_df[orig_stratum_col].to_numpy()
        else:
            orig_strata = np.zeros(n_psus, dtype=np.int64)

        unique_orig, counts = np.unique(orig_strata, return_counts=True)
        stratum_counts = dict(zip(unique_orig.tolist(), counts.tolist()))

        small_strata = [(s, c) for s, c in stratum_counts.items() if c < 2]
        if small_strata:
            raise DimensionError(
                title="Insufficient PSUs per stratum",
                detail=(
                    f"All strata must have at least 2 PSUs. "
                    f"Found {len(small_strata)} strata with fewer."
                ),
                code="INSUFFICIENT_PSU",
                where=where,
                param="stratum",
                expected="≥2 PSUs per stratum",
                got=f"{small_strata[:5]}{'...' if len(small_strata) > 5 else ''}",
                hint="Combine small strata or check design specification.",
            )

        if _method == "brr":
            odd_strata = [(s, c) for s, c in stratum_counts.items() if c % 2 == 1]
            if odd_strata:
                raise DimensionError(
                    title="Odd PSU counts for BRR",
                    detail=(
                        f"BRR requires even number of PSUs per stratum. "
                        f"Found {len(odd_strata)} strata with odd counts."
                    ),
                    code="ODD_PSU_COUNT",
                    where=where,
                    param="stratum",
                    expected="Even PSU count per stratum",
                    got=f"{odd_strata[:5]}{'...' if len(odd_strata) > 5 else ''}",
                    hint="Use method='jk2' which allows 2-3 PSUs per stratum.",
                )

        var_strata = np.empty(n_psus, dtype=np.int64)
        rng = np.random.default_rng(rstate) if shuffle and not order_cols else None
        var_stratum_counter = 0

        for orig_str in unique_orig:
            mask = orig_strata == orig_str
            indices = np.where(mask)[0]
            n_in_stratum = len(indices)

            if rng is not None:
                indices = indices.copy()
                rng.shuffle(indices)

            if method == "brr":
                for i in range(0, n_in_stratum, 2):
                    var_strata[indices[i]] = var_stratum_counter
                    var_strata[indices[i + 1]] = var_stratum_counter
                    var_stratum_counter += 1
            else:  # jk2
                if n_in_stratum % 2 == 1:
                    for i in range(0, n_in_stratum - 3, 2):
                        var_strata[indices[i]] = var_stratum_counter
                        var_strata[indices[i + 1]] = var_stratum_counter
                        var_stratum_counter += 1
                    var_strata[indices[-3]] = var_stratum_counter
                    var_strata[indices[-2]] = var_stratum_counter
                    var_strata[indices[-1]] = var_stratum_counter
                    var_stratum_counter += 1
                else:
                    for i in range(0, n_in_stratum, 2):
                        var_strata[indices[i]] = var_stratum_counter
                        var_strata[indices[i + 1]] = var_stratum_counter
                        var_stratum_counter += 1

        psu_to_var_stratum = dict(zip(psu_list, var_strata))
        psu_vec = df[psu_col].to_list()
        obs_var_strata = np.array([psu_to_var_stratum[p] for p in psu_vec], dtype=np.int64)

        self._sample._data = df.with_columns(pl.Series(name=into, values=obs_var_strata))
        self._sample._design = self._sample._design.update(stratum=into)

        return self._sample

    # ---------- BRR ----------
    def create_brr_wgts(
        self,
        n_reps: int | None = None,
        *,
        rep_prefix: str | None = None,
        fay_coef: float = 0.0,
        rstate: int | None = None,
        drop_nulls: bool = False,
    ) -> Any:
        df = self._sample._data
        design = self._sample._design

        if design.stratum is None:
            raise ValueError("BRR requires `stratum` in Design (got stratum=None).")
        if design.psu is None:
            raise ValueError("BRR requires `psu` in Design (got psu=None).")

        if drop_nulls:
            needed = list(
                {c for c in [design.wgt, design.stratum, design.psu] if isinstance(c, str)}
            )
            data = drop_missing(df=df, cols=needed, treat_infinite_as_missing=True)
        else:
            data = df

        main_weights = _to_float_array(data, design.wgt, len(data))
        stratum_int = _to_int_array(data, design.stratum)
        psu_int = _to_int_array(data, design.psu)

        assert rust_create_brr_wgts is not None  # noqa: S101
        rep_mat, df_val = rust_create_brr_wgts(
            main_weights,
            stratum_int,
            psu_int,
            n_reps,
            fay_coef,
            rstate,
        )
        n_reps_actual = rep_mat.shape[1]

        rep_prefix = rep_prefix or "svy_brr_wgt"
        rep_cols = _name_rep_cols(rep_prefix, n_reps_actual)
        rep_dicts = {col: rep_mat[:, i] for i, col in enumerate(rep_cols)}
        self._sample._data = data.with_columns(
            [pl.Series(name=col, values=vals) for col, vals in rep_dicts.items()]
        )

        self._sample._design = self._sample._design.fill_missing(
            rep_wgts=RepWeights(
                method=EstimationMethod.BRR,
                prefix=rep_prefix,
                n_reps=n_reps_actual,
                fay_coef=fay_coef,
                df=df_val,
            )
        )

        return self._sample

    # ---------- Jackknife ----------
    def create_jk_wgts(
        self,
        *,
        paired: bool = False,
        rep_prefix: str | None = None,
        rstate: int | None = None,
        drop_nulls: bool = False,
    ) -> Any:
        """
        Create Jackknife replicate weights.

        Args:
            paired: If False, use JK1/JKn (delete-one-PSU).
                    If True, use JK2 (paired jackknife).
            rep_prefix: Prefix for replicate weight column names.
            rstate: Random seed for JK2 PSU selection.
            drop_nulls: Drop rows with missing values in design columns.
        """
        df = self._sample._data
        design = self._sample._design

        if design.psu is None:
            raise ValueError("Jackknife requires at least psu; stratum optional.")

        if drop_nulls:
            needed = list(
                {c for c in [design.wgt, design.stratum, design.psu] if isinstance(c, str)}
            )
            data = drop_missing(df=df, cols=needed, treat_infinite_as_missing=True)
        else:
            data = df

        main_weights = _to_float_array(data, design.wgt, len(data))
        psu_int = _to_int_array(data, design.psu)
        stratum_int = _to_int_array(data, design.stratum)

        assert rust_create_jk_wgts is not None  # noqa: S101
        rep_mat, df_val = rust_create_jk_wgts(
            main_weights,
            psu_int,
            stratum_int,
            paired,
            rstate,
        )
        n_reps = rep_mat.shape[1]

        rep_prefix = rep_prefix or ("svy_jk2_wgt" if paired else "svy_jkn_wgt")
        rep_cols = _name_rep_cols(rep_prefix, n_reps)
        rep_dicts = {col: rep_mat[:, i] for i, col in enumerate(rep_cols)}
        self._sample._data = data.with_columns(
            [pl.Series(name=col, values=vals) for col, vals in rep_dicts.items()]
        )

        self._sample._design = self._sample._design.fill_missing(
            rep_wgts=RepWeights(
                method=EstimationMethod.JACKKNIFE,
                prefix=rep_prefix,
                n_reps=n_reps,
                df=df_val,
            )
        )

        return self._sample

    # ---------- Bootstrap ----------
    def create_bs_wgts(
        self,
        n_reps: int = 500,
        *,
        rep_prefix: str | None = None,
        drop_nulls: bool = False,
        rstate: RandomState = None,
    ) -> Any:
        df = self._sample._data
        design = self._sample._design

        if design.psu is None:
            raise ValueError("Bootstrap requires at least psu; stratum optional.")
        if n_reps is None:
            raise ValueError("n_reps must be specified for Bootstrap.")

        if drop_nulls:
            needed = list(
                {c for c in [design.wgt, design.stratum, design.psu] if isinstance(c, str)}
            )
            data = drop_missing(df=df, cols=needed, treat_infinite_as_missing=True)
        else:
            data = df

        main_weights = _to_float_array(data, design.wgt, len(data))
        psu_int = _to_int_array(data, design.psu)
        stratum_int = _to_int_array(data, design.stratum)

        rng = resolve_random_state(rstate)
        seed = (
            int(rng.integers(0, 2**63 - 1))
            if hasattr(rng, "integers")
            else int(rng.randint(0, 2**31 - 1))
        )

        assert rust_create_bootstrap_wgts is not None  # noqa: S101
        rep_mat, df_val = rust_create_bootstrap_wgts(
            main_weights,
            psu_int,
            n_reps,
            stratum_int,
            seed,
        )

        rep_prefix = rep_prefix or "svy_boot_wgt"
        rep_cols = _name_rep_cols(rep_prefix, n_reps)
        rep_dicts = {col: rep_mat[:, i] for i, col in enumerate(rep_cols)}
        self._sample._data = data.with_columns(
            [pl.Series(name=col, values=vals) for col, vals in rep_dicts.items()]
        )

        self._sample._design = self._sample._design.update_rep_weights(
            method=EstimationMethod.BOOTSTRAP,
            prefix=rep_prefix,
            n_reps=n_reps,
            df=df_val,
        )

        return self._sample

    # ---------- SDR ----------
    def create_sdr_wgts(
        self,
        n_reps: int = 4,
        *,
        rep_prefix: str | None = None,
        order_col: str | None = None,
        drop_nulls: bool = False,
    ) -> Any:
        """
        Create SDR (Successive Difference Replication) weights.

        Args:
            n_reps: Number of replicates (typically 4 or more).
            rep_prefix: Prefix for replicate weight column names.
            order_col: Column specifying sort order within strata.
            drop_nulls: Drop rows with missing values in design columns.
        """
        df = self._sample._data
        design = self._sample._design

        if n_reps < 2:
            raise ValueError("SDR requires at least 2 replicates.")

        if drop_nulls:
            needed = list({c for c in [design.wgt, design.stratum] if isinstance(c, str)})
            if order_col:
                needed.append(order_col)
            data = drop_missing(df=df, cols=needed, treat_infinite_as_missing=True)
        else:
            data = df

        main_weights = _to_float_array(data, design.wgt, len(data))
        stratum_int = _to_int_array(data, design.stratum)

        # Order array — single column only
        order_int: np.ndarray | None = None
        if order_col and order_col in data.columns:
            order_int = _to_int_array(data, order_col)

        assert rust_create_sdr_wgts is not None  # noqa: S101
        rep_mat, df_val = rust_create_sdr_wgts(
            main_weights,
            n_reps,
            stratum_int,
            order_int,
        )

        rep_prefix = rep_prefix or "svy_sdr_wgt"
        rep_cols = _name_rep_cols(rep_prefix, n_reps)
        rep_dicts = {col: rep_mat[:, i] for i, col in enumerate(rep_cols)}
        self._sample._data = data.with_columns(
            [pl.Series(name=col, values=vals) for col, vals in rep_dicts.items()]
        )

        self._sample._design = self._sample._design.update_rep_weights(
            method=EstimationMethod.SDR,
            prefix=rep_prefix,
            n_reps=n_reps,
            df=df_val,
        )

        return self._sample

    # ------------------------------------------------------------------ #
    # Reweighting: Adjustment (Non-Response / Within-Class)
    # ------------------------------------------------------------------ #

    def adjust(
        self,
        resp_status: str,
        by: Feature | None,
        *,
        resp_mapping: DomainScalarMap | None = None,
        wgt_name: str | None = None,
        replace: bool = False,
        rep_wgts_prefix: str | None = None,
        ignore_reps: bool = False,
        unknown_to_inelig: bool = True,
        update_design_wgts: bool = True,
        keep_resp: bool = True,
    ) -> Any:
        df = self._sample._data
        design = self._sample._design

        if design.wgt is None:
            raise TypeError("Sample weight is None! Set design.wgt before calling adjust().")
        wgt = design.wgt
        if wgt not in df.columns:
            raise KeyError(f"Weight column {wgt!r} not found in data.")
        if not isinstance(resp_status, str) or resp_status not in df.columns:
            raise TypeError(
                f"`resp_status` must be a string representing an existing column name. "
                f"Got: {resp_status!r}"
            )
        if replace and wgt_name is not None:
            raise TypeError("Pass either replace=True or wgt_name=..., not both.")

        _, adj_class_arr = _normalize_by_term(
            df, by, param_name="by", where="Sample.weighting.adjust"
        )

        wgt_arr = df.get_column(wgt).to_numpy()
        resp_status_arr = df.get_column(resp_status).to_numpy()

        adj_wgt_arr = _adjust_nr(
            wgts=wgt_arr,
            adj_class=adj_class_arr,
            resp_status=resp_status_arr,
            resp_mapping=resp_mapping,
            unknown_to_inelig=unknown_to_inelig,
        )

        existing_cols = set(df.columns)

        if replace:
            target_wgt = wgt
            df = df.with_columns(pl.Series(target_wgt, adj_wgt_arr))
        else:
            target_wgt = _unique_name(
                f"{SVY_ADJ_PREFIX}{wgt}" if wgt_name is None else wgt_name, existing_cols
            )
            df = df.with_columns(pl.Series(target_wgt, adj_wgt_arr))
            existing_cols.add(target_wgt)

        if update_design_wgts:
            self._sample._design = self._sample._design.update(wgt=target_wgt)

        if not ignore_reps and design.rep_wgts is not None:
            rep_cols = design.rep_wgts.columns
            wgts_arr = df.select(rep_cols).to_numpy()

            if adj_class_arr is not None:
                unique_classes = np.unique(adj_class_arr)
                class_to_idx = {cls: idx for idx, cls in enumerate(unique_classes)}
                adj_class_indices = np.array(
                    [class_to_idx[cls] for cls in adj_class_arr], dtype=np.int64
                )
            else:
                adj_class_indices = np.zeros(len(wgt_arr), dtype=np.int64)

            resp_codes = np.zeros(len(resp_status_arr), dtype=np.int64)
            resp_lower = np.char.lower(resp_status_arr.astype(str))
            CANONICAL_TO_INT = {"rr": 0, "nr": 1, "in": 2, "uk": 3}

            if resp_mapping is not None:
                for canonical_key, data_label in resp_mapping.items():
                    key_lower = str(canonical_key).lower()
                    if key_lower not in CANONICAL_TO_INT:
                        raise ValueError(
                            f"Invalid canonical key '{canonical_key}'. "
                            f"Must be one of: rr, nr, in, uk"
                        )
                    int_code = CANONICAL_TO_INT[key_lower]
                    resp_codes[resp_lower == str(data_label).lower()] = int_code
            else:
                resp_codes[resp_lower == "rr"] = 0
                resp_codes[resp_lower == "nr"] = 1
                resp_codes[resp_lower == "in"] = 2
                resp_codes[resp_lower == "uk"] = 3

            assert rust_adjust_nr is not None  # noqa: S101
            adj_wgts_arr = rust_adjust_nr(
                wgts_arr,
                adj_class_indices,
                resp_codes,
                unknown_to_inelig,
            )

            wgts_df = pl.DataFrame(adj_wgts_arr)
            n_reps = len(rep_cols)

            if rep_wgts_prefix is not None:
                final_prefix = rep_wgts_prefix
            else:
                final_prefix = f"{SVY_ADJ_PREFIX}{design.rep_wgts.prefix}"

            new_rep_names = [
                _unique_name(f"{final_prefix}{i}", existing_cols) for i in range(1, n_reps + 1)
            ]
            wgts_df.columns = new_rep_names

            self._sample._data = df.hstack(wgts_df)
            df = self._sample._data

            if update_design_wgts:
                self._sample._design = self._sample._design.update_rep_weights(
                    method=design.rep_wgts.method,
                    prefix=final_prefix,
                    n_reps=n_reps,
                    fay_coef=design.rep_wgts.fay_coef,
                    df=design.rep_wgts.df,
                )

        self._sample._data = df

        if keep_resp:

            def _resp_mask_expr(col: str, mapping: dict | None) -> pl.Expr:
                if mapping is None:
                    return pl.col(col) == "rr"
                rr_codes = mapping.get("rr", "rr")
                if isinstance(rr_codes, (list, tuple, set, np.ndarray)):
                    return pl.col(col).is_in(list(rr_codes))
                else:
                    return pl.col(col) == rr_codes

            mask = _resp_mask_expr(resp_status, resp_mapping)
            self._sample._data = self._sample._data.filter(mask)

        return self._sample

    # ------------------------------------------------------------------ #
    # Reweighting: Normalization
    # ------------------------------------------------------------------ #

    def normalize(
        self,
        controls: DomainScalarMap | Number | None = None,
        *,
        by: Feature | None = None,
        wgt_name: str | None = None,
        replace: bool = False,
        rep_wgts_prefix: str | None = None,
        ignore_reps: bool = False,
        update_design_wgts: bool = True,
    ) -> Any:
        df = self._sample._data
        design = self._sample._design

        if design.wgt is None:
            raise TypeError("Sample weight is None! Set design.wgt before calling normalize().")
        wgt = design.wgt
        if wgt not in df.columns:
            raise KeyError(f"Weight column {wgt!r} not found in data.")
        if replace and wgt_name is not None:
            raise TypeError("Pass either replace=True or wgt_name=..., not both.")

        _, by_arr = _normalize_by_term(df, by, param_name="by", where="Sample.weighting.normalize")

        wgt_arr = df.get_column(wgt).to_numpy()
        existing_cols = set(df.columns)

        samp_norm, _, _ = _normalize(wgt=wgt_arr, control=controls, by_arr=by_arr)

        if replace:
            target_wgt = wgt
            df = df.with_columns(pl.Series(name=target_wgt, values=samp_norm))
        else:
            target_wgt = _unique_name(wgt_name or f"{SVY_NORM_PREFIX}{wgt}", existing_cols)
            df = df.with_columns(pl.Series(name=target_wgt, values=samp_norm))
            existing_cols.add(target_wgt)

        if update_design_wgts:
            self._sample._design = self._sample._design.update(wgt=target_wgt)

        if not ignore_reps and design.rep_wgts is not None:
            rep_cols = design.rep_wgts.columns

            if rep_cols:
                wgts_arr = df.select(rep_cols).to_numpy()

                if by_arr is None:
                    assert rust_normalize is not None  # noqa: S101
                    adj_wgts_arr = rust_normalize(wgts_arr, None, None)
                else:
                    raise NotImplementedError(
                        "Grouped normalization on replicate weights not yet supported. "
                        "Use by=None for global normalization."
                    )

                if rep_wgts_prefix is None:
                    final_prefix = f"{SVY_NORM_PREFIX}{design.rep_wgts.prefix}"
                else:
                    final_prefix = rep_wgts_prefix

                n_reps = len(rep_cols)
                new_rep_names = [f"{final_prefix}{i}" for i in range(1, n_reps + 1)]

                wgts_df = pl.DataFrame(adj_wgts_arr, schema=new_rep_names)
                self._sample._data = df.hstack(wgts_df)
                df = self._sample._data

                if update_design_wgts:
                    self._sample._design = self._sample._design.update_rep_weights(
                        method=design.rep_wgts.method,
                        prefix=final_prefix,
                        n_reps=n_reps,
                        fay_coef=design.rep_wgts.fay_coef,
                        df=design.rep_wgts.df,
                    )

        self._sample._data = df
        return self._sample

    # ------------------------------------------------------------------ #
    # Reweighting: Poststratification
    # ------------------------------------------------------------------ #

    def poststratify(
        self,
        controls: DomainScalarMap | Number | None = None,
        *,
        factors: DomainScalarMap | Number | None = None,
        by: Feature | None = None,
        wgt_name: str | None = None,
        replace: bool = False,
        rep_wgts_prefix: str | None = None,
        ignore_reps: bool = False,
        update_design_wgts: bool = True,
    ) -> Any:
        if controls is None and factors is None:
            raise AssertionError("control or factor must be specified.")

        df = self._sample._data
        design = self._sample._design

        if design.wgt is None:
            raise TypeError("Sample weight is None! Set design.wgt before calling poststratify().")
        wgt = design.wgt
        if wgt not in df.columns:
            raise KeyError(f"Weight column {wgt!r} not found in data.")
        if replace and wgt_name is not None:
            raise ValueError("Pass either replace=True or wgt_name=..., not both.")

        _, by_arr = _normalize_by_term(
            df, by, param_name="by", where="Sample.weighting.poststratify"
        )

        wgt_arr = df.get_column(wgt).to_numpy()
        existing_cols = set(df.columns)

        ps_wgt = _poststratify(
            wgt=wgt_arr,
            control=controls,
            factor=factors,
            by_arr=by_arr,
        )

        if replace:
            target_wgt = wgt
            df = df.with_columns(pl.Series(name=target_wgt, values=ps_wgt))
        else:
            target_wgt = _unique_name(wgt_name or f"{SVY_PS_PREFIX}{wgt}", existing_cols)
            df = df.with_columns(pl.Series(name=target_wgt, values=ps_wgt))
            existing_cols.add(target_wgt)

        if update_design_wgts:
            self._sample._design = self._sample._design.update(wgt=target_wgt)

        if not ignore_reps and design.rep_wgts is not None:
            rep_cols = design.rep_wgts.columns

            if rep_cols:
                wgts_arr = df.select(rep_cols).to_numpy()

                if by_arr is None:
                    by_indices = np.zeros(len(wgt_arr), dtype=np.int64)

                    if isinstance(controls, dict):
                        if len(controls) != 1:
                            raise ValueError(
                                "When by=None (single stratum), controls dict must have "
                                "exactly one key"
                            )
                        control_arr = np.array(
                            [float(next(iter(controls.values())))], dtype=np.float64
                        )
                    elif isinstance(controls, (int, float)):
                        control_arr = np.array([float(controls)], dtype=np.float64)
                    elif factors is not None:
                        if isinstance(factors, dict):
                            if len(factors) != 1:
                                raise ValueError(
                                    "When by=None (single stratum), factors dict must have "
                                    "exactly one key"
                                )
                            factor_val = float(next(iter(factors.values())))
                            control_arr = np.array([factor_val * wgt_arr.sum()], dtype=np.float64)
                        else:
                            control_arr = np.array(
                                [float(factors) * wgt_arr.sum()], dtype=np.float64
                            )
                    else:
                        raise ValueError(
                            "Must provide controls or factors for post-stratification"
                        )

                    assert rust_poststratify is not None  # noqa: S101
                    adj_wgts_arr = rust_poststratify(wgts_arr, by_indices, control_arr)

                elif isinstance(controls, dict):
                    unique_vals = np.unique(by_arr)
                    val_to_idx = {val: idx for idx, val in enumerate(unique_vals)}
                    by_indices = np.array([val_to_idx[val] for val in by_arr], dtype=np.int64)
                    control_arr = np.array(
                        [float(controls[val]) for val in unique_vals], dtype=np.float64
                    )

                    assert rust_poststratify is not None  # noqa: S101
                    adj_wgts_arr = rust_poststratify(wgts_arr, by_indices, control_arr)

                elif isinstance(controls, (int, float)):
                    unique_vals = np.unique(by_arr)
                    val_to_idx = {val: idx for idx, val in enumerate(unique_vals)}
                    by_indices = np.array([val_to_idx[val] for val in by_arr], dtype=np.int64)
                    n_strata = len(unique_vals)
                    control_arr = np.zeros(n_strata, dtype=np.float64)
                    total_weight = wgt_arr.sum()
                    for idx, val in enumerate(unique_vals):
                        mask = by_arr == val
                        stratum_weight = wgt_arr[mask].sum()
                        control_arr[idx] = float(controls) * (stratum_weight / total_weight)

                    assert rust_poststratify is not None  # noqa: S101
                    adj_wgts_arr = rust_poststratify(wgts_arr, by_indices, control_arr)

                elif factors is not None:
                    unique_vals = np.unique(by_arr)
                    val_to_idx = {val: idx for idx, val in enumerate(unique_vals)}
                    by_indices = np.array([val_to_idx[val] for val in by_arr], dtype=np.int64)

                    if isinstance(factors, dict):
                        control_arr = np.zeros(len(unique_vals), dtype=np.float64)
                        for idx, val in enumerate(unique_vals):
                            mask = by_arr == val
                            stratum_weight = wgt_arr[mask].sum()
                            control_arr[idx] = float(factors[val]) * stratum_weight
                    else:
                        control_arr = np.zeros(len(unique_vals), dtype=np.float64)
                        for idx, val in enumerate(unique_vals):
                            mask = by_arr == val
                            stratum_weight = wgt_arr[mask].sum()
                            control_arr[idx] = float(factors) * stratum_weight

                    assert rust_poststratify is not None  # noqa: S101
                    adj_wgts_arr = rust_poststratify(wgts_arr, by_indices, control_arr)

                else:
                    raise ValueError("Must provide controls or factors for post-stratification")

                if rep_wgts_prefix is None:
                    final_prefix = f"{SVY_PS_PREFIX}{design.rep_wgts.prefix}"
                else:
                    final_prefix = rep_wgts_prefix

                n_reps = len(rep_cols)
                new_rep_names = [f"{final_prefix}{i}" for i in range(1, n_reps + 1)]

                wgts_df = pl.DataFrame(adj_wgts_arr, schema=new_rep_names)
                self._sample._data = df.hstack(wgts_df)
                df = self._sample._data

                if update_design_wgts:
                    self._sample._design = self._sample._design.update_rep_weights(
                        method=design.rep_wgts.method,
                        prefix=final_prefix,
                        n_reps=n_reps,
                        fay_coef=design.rep_wgts.fay_coef,
                        df=design.rep_wgts.df,
                    )

        self._sample._data = df
        return self._sample

    # ------------------------------------------------------------------ #
    # Raking
    # ------------------------------------------------------------------ #

    def controls_margins_template(
        self,
        *,
        margins: Mapping[str, str],
        cat_na: str = "level",
        na_label: str = "__NA__",
    ) -> dict[str, dict[Category, float]]:
        df: pl.DataFrame = self._sample.data
        where = "Sample.weighting.controls_margins_template"

        for mname, col in margins.items():
            if not isinstance(mname, str) or not isinstance(col, str):
                raise TypeError(
                    f"{where}: margins must be {{str -> str}}; got {type(mname)!r}->{type(col)!r}"
                )
            if col not in df.columns:
                raise KeyError(f"{where}: column {col!r} (for margin {mname!r}) not found.")

        result: dict[str, dict[Category, float]] = {}

        for mname, col in margins.items():
            s = df.get_column(col)

            if cat_na not in ("error", "level"):
                raise ValueError(f"{where}: cat_na must be 'error' or 'level'.")

            if cat_na == "error":
                if s.is_null().any():
                    raise DimensionError(
                        title="Missing values in margin column",
                        detail=f"Nulls found in {col!r}. Choose cat_na='level' or fix data.",
                        code="MARGIN_NA",
                        where=where,
                        param=col,
                        hint="Use cat_na='level' to include a missing category.",
                    )
                s_norm = s.cast(pl.Utf8)
            else:
                s_norm = s.cast(pl.Utf8).fill_null(na_label)

            cats = pl.Series("__cats__", s_norm.unique().to_list(), dtype=pl.Utf8).to_list()
            cats_sorted = sorted(cats, key=_num_sort_key_label)
            result[mname] = {lab: np.nan for lab in cats_sorted}

        return result

    def rake(
        self,
        *,
        controls: ControlsType | None = None,
        factors: ControlsType | None = None,
        wgt_name: str | None = None,
        replace: bool = False,
        rep_wgts_prefix: str | None = None,
        ignore_reps: bool = False,
        ll_bound: float | None = None,
        up_bound: float | None = None,
        tol: float = 1e-4,
        max_iter: int = 100,
        display_iter: bool = False,
        update_design_wgts: bool = True,
    ) -> "Sample":
        df = self._sample._data
        design = self._sample._design

        if design.wgt is None:
            raise TypeError("Sample weight is None! Set design.wgt before calling rake().")
        wgt = design.wgt
        if wgt not in df.columns:
            raise KeyError(f"Weight column {wgt!r} not found in data.")
        if replace and wgt_name is not None:
            raise TypeError("Pass either replace=True or wgt_name=..., not both.")

        controls_norm: ControlsType | None = _normalize_controls_like(x=controls)
        factors_norm: ControlsType | None = _normalize_controls_like(x=factors)

        if controls_norm is None and factors_norm is None:
            raise ValueError("Either 'control' or 'factor' must be specified.")
        if controls_norm is not None and factors_norm is not None:
            raise ValueError("Provide exactly one of 'control' or 'factor'.")

        if ll_bound is not None and up_bound is not None and ll_bound > up_bound:
            raise ValueError("ll_bound cannot be greater than up_bound.")

        rake_cols = (
            list(controls_norm.keys()) if controls_norm is not None else list(factors_norm.keys())  # type: ignore[union-attr]
        )
        if not rake_cols:
            raise ValueError("No raking columns provided in control/factor keys.")

        processed: dict[str, np.ndarray] = {}
        w0 = df.get_column(wgt).to_numpy()

        for col in rake_cols:
            if not isinstance(col, str) or col not in df.columns:
                raise KeyError(f"Raking column {col!r} not found in data.")
            s = df.get_column(col)
            if s.len() != w0.size:
                raise ValueError(f"Raking column {col!r} length mismatch.")
            if s.null_count() > 0:
                raise ValueError(f"Raking column {col!r} contains nulls.")
            processed[col] = s.to_numpy()

        control_final: ControlsType = controls_norm or _calculate_controls_from_factors(
            wgts=w0,
            margins=processed,
            factors=cast(ControlsType, factors_norm),
        )

        missing = [m for m in processed if m not in control_final]
        extra = [m for m in control_final if m not in processed]
        if missing or extra:
            raise ValueError(
                "Control totals must be provided for all specified raking columns "
                f"(missing={missing or []}, extra={extra or []})."
            )

        for col_name, totals in control_final.items():
            if not isinstance(totals, Mapping) or not totals:
                raise ValueError(f"Control for column {col_name!r} must be a non-empty mapping.")
            vals = np.array(list(totals.values()), dtype=float)
            if not np.all(np.isfinite(vals)) or np.any(vals < 0):
                raise ValueError(
                    f"Control totals for {col_name!r} must be finite and non-negative."
                )
            if np.all(vals == 0):
                raise ValueError(f"All-zero control totals for {col_name!r} not allowed.")

        raked_w = _rake(
            wgt=w0,
            margins=processed,
            control=control_final,
            ll_bound=ll_bound,
            up_bound=up_bound,
            tol=tol,
            max_iter=max_iter,
            display_iter=display_iter,
        )

        existing_cols = set(df.columns)
        if replace:
            target_wgt = wgt
            df = df.with_columns(pl.Series(name=target_wgt, values=raked_w))
        else:
            target_wgt = _unique_name(wgt_name or f"{SVY_RAKED_PREFIX}{wgt}", existing_cols)
            df = df.with_columns(pl.Series(name=target_wgt, values=raked_w))
            existing_cols.add(target_wgt)

        if update_design_wgts:
            self._sample._design = self._sample._design.update(wgt=target_wgt)

        if not ignore_reps and design.rep_wgts is not None:
            rep_cols = design.rep_wgts.columns

            if rep_cols:
                if rep_wgts_prefix is None:
                    final_prefix = f"{SVY_RAKED_PREFIX}{design.rep_wgts.prefix}"
                else:
                    final_prefix = rep_wgts_prefix

                n_reps = len(rep_cols)
                wgts_arr = df.select(rep_cols).to_numpy()

                margin_indices = []
                margin_targets = []

                for col in rake_cols:
                    cats_sorted = sorted(control_final[col].keys(), key=_num_sort_key_label)
                    cat_to_idx = {cat: idx for idx, cat in enumerate(cats_sorted)}
                    indices = np.array([cat_to_idx[val] for val in processed[col]], dtype=np.int64)
                    targets = np.array(
                        [float(control_final[col][cat]) for cat in cats_sorted],
                        dtype=np.float64,
                    )
                    margin_indices.append(indices)
                    margin_targets.append(targets)

                assert rust_rake is not None  # noqa: S101
                raked_reps = rust_rake(
                    wgts_arr,
                    margin_indices,
                    margin_targets,
                    ll_bound,
                    up_bound,
                    tol,
                    max_iter,
                )

                new_rep_names = [f"{final_prefix}{i}" for i in range(1, n_reps + 1)]
                wgts_df = pl.DataFrame(raked_reps, schema=new_rep_names)
                self._sample._data = df.hstack(wgts_df)
                df = self._sample._data

                if update_design_wgts:
                    self._sample._design = self._sample._design.update_rep_weights(
                        method=design.rep_wgts.method,
                        prefix=final_prefix,
                        n_reps=n_reps,
                        fay_coef=design.rep_wgts.fay_coef,
                        df=design.rep_wgts.df,
                    )

        self._sample._data = df
        return self._sample

    # ------------------------------------------------------------------ #
    # Calibration
    # ------------------------------------------------------------------ #

    def control_aux_template(
        self,
        *,
        x: Sequence[Feature],
        by: Feature | None = None,
        by_na: Literal["error", "level", "drop"] = "error",
        na_label: str = "__NA__",
    ) -> dict[Category, Number] | dict[Category, dict[Category, Number]]:
        _, shape = self.build_aux_matrix(x=x, by=by, by_na=by_na, na_label=na_label)
        return shape

    def build_aux_matrix(
        self,
        *,
        x: Sequence[Feature],
        by: Feature | None = None,
        by_na: Literal["error", "level", "drop"] = "error",
        na_label: str = "__NA__",
    ) -> tuple[np.ndarray, dict[Category, Number] | dict[Category, dict[Category, Number]]]:
        df: pl.DataFrame = self._sample.data
        where = "Sample.weighting.build_aux_matrix"

        by_n, _ = _normalize_by_term(df, by, param_name="by", where=where)

        if not x:
            raise MethodError.not_applicable(
                where=where, method="aux builder", reason="No terms specified."
            )

        all_exprs = []
        all_labels = []

        for term in x:
            t_exprs, t_labs = _expand_term(term, df, where)
            all_exprs.extend(t_exprs)
            all_labels.extend(t_labs)

        if not all_exprs:
            raise MethodError.not_applicable(
                where=where, method="aux builder", reason="Terms produced no columns."
            )

        final_exprs = [e.alias(f"__aux_col_{i}") for i, e in enumerate(all_exprs)]

        if by_n is None:
            X_df = df.select(final_exprs)
            X = X_df.to_numpy()
            return X, {lab: np.nan for lab in all_labels}

        if by_na == "error":
            if (
                df.select([pl.col(c).is_null().any() for c in by_n])
                .select(pl.any_horizontal(pl.all()))
                .item()
            ):
                raise DimensionError(
                    title="Missing values in `by` columns",
                    detail=f"Nulls found in {by_n}.",
                    code="BY_NA",
                    where=where,
                    param="by",
                    hint="Use by_na='level' or 'drop', or fix data.",
                )
            by_exprs = [pl.col(c) for c in by_n]
        elif by_na == "level":
            by_exprs = [pl.col(c).fill_null(na_label) for c in by_n]
        elif by_na == "drop":
            filter_expr = pl.all_horizontal([pl.col(c).is_not_null() for c in by_n])
            df = df.filter(filter_expr)
            by_exprs = [pl.col(c) for c in by_n]
        else:
            raise ValueError("by_na must be 'error', 'level', or 'drop'.")

        combined_df = df.select(final_exprs + by_exprs)
        n_x_cols = len(final_exprs)
        X = combined_df[:, :n_x_cols].to_numpy()
        by_arr = combined_df[:, n_x_cols:].to_numpy()

        _keys_raw = [row[0] if by_arr.shape[1] == 1 else tuple(row.tolist()) for row in by_arr]
        keys: list[Category] = _keys_raw  # type: ignore[assignment]

        try:
            uniq_keys = sorted(list(set(keys)))
        except TypeError:
            uniq_keys = list(set(keys))

        inner_template = {lab: np.nan for lab in all_labels}
        return X, {k: inner_template.copy() for k in uniq_keys}

    def calibrate(
        self,
        *,
        controls: dict[Feature, Any],
        by: Feature | None = None,
        scale: Number | list[Number] | np.ndarray = 1.0,
        bounded: bool = False,
        additive: bool = False,
        wgt_name: str | None = None,
        replace: bool = False,
        update_design_wgts: bool = True,
        rep_wgts_prefix: str | None = None,
        ignore_reps: bool = False,
    ) -> Any:
        where = "Sample.weighting.calibrate"

        if not isinstance(controls, dict) or not controls:
            raise MethodError.not_applicable(
                where=where,
                method="calibrate",
                reason="`controls` must be a non-empty dictionary.",
                param="controls",
            )

        terms: list[Feature]
        is_global = by is None

        if is_global:
            terms = list(controls.keys())
        else:
            first_domain_val = next(iter(controls.values()))
            if not isinstance(first_domain_val, dict):
                raise TypeError(
                    f"When `by` is used, `controls` values must be dictionaries of "
                    f"targets. Got {type(first_domain_val)}."
                )
            terms = list(first_domain_val.keys())

        X, shape_template = self.build_aux_matrix(x=terms, by=by)

        if is_global:
            x_labels = list(shape_template.keys())
        else:
            x_labels = list(next(iter(shape_template.values())).keys())

        final_control_arg: Any

        if is_global:
            flat_targets = []
            cursor = 0
            for term in terms:
                _, term_labs = _expand_term(term, self._sample.data, where)
                k = len(term_labs)
                segment_labels = x_labels[cursor : cursor + k]
                if segment_labels != term_labs:
                    raise RuntimeError("Internal error: Design matrix label alignment mismatch.")
                spec = controls[term]
                vals = _match_term_targets(segment_labels, spec, str(term))
                flat_targets.extend(vals)
                cursor += k
            final_control_arg = np.array(flat_targets, dtype=float)
        else:
            final_control_arg = {}
            expected_domains = set(shape_template.keys())
            provided_domains = set(controls.keys())
            missing = expected_domains - provided_domains
            extra = provided_domains - expected_domains
            if missing or extra:
                raise ValueError(f"Domain mismatch in controls. Missing={missing}, Extra={extra}")

            for domain, domain_specs in controls.items():
                flat_targets = []
                cursor = 0
                for term in terms:
                    if term not in domain_specs:
                        raise ValueError(
                            f"Missing specification for term '{term}' in domain '{domain}'."
                        )
                    _, term_labs = _expand_term(term, self._sample.data, where)
                    k = len(term_labs)
                    segment_labels = x_labels[cursor : cursor + k]
                    spec = domain_specs[term]
                    vals = _match_term_targets(segment_labels, spec, str(term))
                    flat_targets.extend(vals)
                    cursor += k
                final_control_arg[domain] = dict(zip(x_labels, flat_targets))

        return self.calibrate_matrix(
            aux_vars=X,
            control=final_control_arg,
            by=by,
            scale=scale,
            bounded=bounded,
            additive=additive,
            wgt_name=wgt_name,
            replace=replace,
            update_design_wgts=update_design_wgts,
            labels=x_labels,
            weights_only=False,
            rep_wgts_prefix=rep_wgts_prefix,
            ignore_reps=ignore_reps,
        )

    def calibrate_matrix(
        self,
        *,
        aux_vars: np.ndarray,
        control: Any,
        by: Feature | None = None,
        scale: Number | Sequence[Number] | np.ndarray = 1.0,
        wgt_name: str | None = None,
        replace: bool = False,
        update_design_wgts: bool = True,
        labels: Sequence[Category] | None = None,
        weights_only: bool = False,
        bounded: bool = False,
        additive: bool = False,
        rep_wgts_prefix: str | None = None,
        ignore_reps: bool = False,
    ) -> Any:
        where = "Sample.weighting.calibrate_matrix"
        df: pl.DataFrame = self._sample.data
        design = self._sample._design

        wgt_col = getattr(self._sample.design, "wgt")
        if not isinstance(wgt_col, str) or wgt_col not in df.columns:
            raise MethodError.invalid_choice(
                where=where,
                param="design.wgt",
                got=wgt_col,
                allowed=df.columns,
                hint="Design must reference an existing weight column.",
            )
        w = df.get_column(wgt_col).to_numpy()

        X = np.asarray(aux_vars, dtype=float)
        if X.ndim == 1:
            X = X.reshape(-1, 1)
        if X.shape[0] != df.height:
            raise DimensionError(
                title="Shape mismatch",
                detail=f"aux_vars has {X.shape[0]} rows; data has {df.height}.",
                code="SHAPE_MISMATCH",
                where=where,
                param="aux_vars",
                expected=f"(n,{X.shape[1]}) with n={df.height}",
                got=f"(n={X.shape[0]},{X.shape[1]})",
            )

        _, domain_vec = _normalize_by_term(df, by, param_name="by", where=where)

        if domain_vec is None and isinstance(control, dict) and labels is not None:
            missing_labels = [lbl for lbl in labels if lbl not in control]
            extra_keys = [k for k in control.keys() if k not in labels]
            if missing_labels or extra_keys:
                raise ValueError(
                    f"Control keys mismatch. Missing: {missing_labels}, Extra: {extra_keys}"
                )
            control = np.array([float(control[lbl]) for lbl in labels])

        new_w = _calibrate(
            samp_weight=w,
            aux_vars=X,
            control=control,
            domain=domain_vec,
            scale=scale,
            bounded=bounded,
            additive=additive,
        )

        if weights_only:
            return new_w

        existing_cols = set(df.columns)

        wgt_col_name = wgt_name or f"{SVY_CALIB_PREFIX}{wgt_col}"

        if wgt_col_name in df.columns:
            raise MethodError.invalid_choice(
                where=where,
                param="wgt_name",
                got=wgt_col_name,
                allowed=[c for c in df.columns if c != wgt_col_name],
                hint="Choose a new column name or drop the existing one first.",
            )

        df = df.with_columns(pl.Series(name=wgt_col_name, values=new_w))
        existing_cols.add(wgt_col_name)

        if replace and update_design_wgts:
            if hasattr(design, "update") and callable(getattr(design, "update")):
                self._sample._design = design.update(wgt=wgt_col_name)
            else:
                setattr(design, "wgt", wgt_col_name)

        if not ignore_reps and design.rep_wgts is not None:
            rep_cols = design.rep_wgts.columns

            if rep_cols:
                wgts_arr = df.select(rep_cols).to_numpy()

                if isinstance(scale, (int, float)):
                    scale_arr = np.full(len(w), float(scale), dtype=np.float64)
                else:
                    scale_arr = np.asarray(scale, dtype=np.float64)

                if domain_vec is None:
                    if isinstance(control, dict):
                        if labels is not None:
                            totals_arr = np.array(
                                [float(control[lbl]) for lbl in labels],  # type: ignore[index]
                                dtype=np.float64,
                            )
                        else:
                            totals_arr = np.array(list(control.values()), dtype=np.float64)
                    else:
                        totals_arr = np.asarray(control, dtype=np.float64)

                    n_reps = len(rep_cols)
                    if n_reps >= 10:
                        assert rust_calibrate_parallel is not None  # noqa: S101
                        calib_reps = rust_calibrate_parallel(
                            wgts_arr,
                            X,
                            totals_arr,
                            scale_arr,
                        )
                    else:
                        assert rust_calibrate is not None  # noqa: S101
                        calib_reps = rust_calibrate(
                            wgts_arr,
                            X,
                            totals_arr,
                            scale_arr,
                            additive,
                        )
                else:
                    unique_domains = np.unique(domain_vec)
                    domain_to_idx = {d: idx for idx, d in enumerate(unique_domains)}
                    domain_indices = np.array(
                        [domain_to_idx[d] for d in domain_vec], dtype=np.int64
                    )

                    controls_dict: dict[int, list[float]] = {}
                    for domain in unique_domains:
                        domain_idx = domain_to_idx[domain]
                        if isinstance(control, dict):
                            domain_control = control[domain]
                            if isinstance(domain_control, dict):
                                if labels is not None:
                                    totals = [float(domain_control[lbl]) for lbl in labels]
                                else:
                                    totals = list(domain_control.values())
                            else:
                                totals = list(np.asarray(domain_control, dtype=float))
                        else:
                            totals = list(np.asarray(control, dtype=float))
                        controls_dict[domain_idx] = totals

                    assert rust_calibrate_by_domain is not None  # noqa: S101
                    calib_reps = rust_calibrate_by_domain(
                        wgts_arr,
                        X,
                        domain_indices,
                        controls_dict,
                        scale_arr,
                        additive,
                    )

                if rep_wgts_prefix is None:
                    final_prefix = f"{SVY_CALIB_PREFIX}{design.rep_wgts.prefix}"
                else:
                    final_prefix = rep_wgts_prefix

                n_reps = len(rep_cols)
                new_rep_names = [f"{final_prefix}{i}" for i in range(1, n_reps + 1)]

                wgts_df = pl.DataFrame(calib_reps, schema=new_rep_names)
                df = df.hstack(wgts_df)

                if update_design_wgts:
                    self._sample._design = self._sample._design.update_rep_weights(
                        method=design.rep_wgts.method,
                        prefix=final_prefix,
                        n_reps=n_reps,
                        fay_coef=design.rep_wgts.fay_coef,
                        df=design.rep_wgts.df,
                    )

        self._sample._data = df
        return self._sample
