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

import logging
import math
import re

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

import numpy as np
import polars as pl
import svy_rs as rs

from scipy import stats

from svy.core.constants import _BY_SEP, _INTERNAL_CONCAT_SUFFIX
from svy.core.enumerations import EstimationMethod, PopParam, QuantileMethod
from svy.core.types import ExprLike, WhereArg
from svy.estimation.estimate import Estimate, ParamEst
from svy.utils.helpers import _colspec_to_list


log = logging.getLogger(__name__)

if TYPE_CHECKING:
    from svy.core.sample import Sample

# Domain estimation constants
_DOMAIN_COL = "_svy_domain_"


class Estimation:
    def __init__(self, sample: Sample) -> None:
        self._sample = sample
        self._design_cache: dict[str, Any] | None = None
        self._polars_cache: dict[str, Any] | None = None

    def _get_factorized_design(self) -> dict[str, Any]:
        if self._design_cache is not None:
            _d = self._sample._data
            _h = (
                cast(pl.DataFrame, _d.collect()).height
                if isinstance(_d, pl.LazyFrame)
                else cast(pl.DataFrame, _d).height
            )
            if self._design_cache["n_rows"] == _h:
                return self._design_cache
            else:
                self._design_cache = None

        _raw_data = self._sample._data
        local_data: pl.DataFrame = (
            cast(pl.DataFrame, _raw_data.collect())
            if isinstance(_raw_data, pl.LazyFrame)
            else cast(pl.DataFrame, _raw_data)
        )
        design = self._sample._design

        cache: dict[str, Any] = {
            "n_rows": local_data.height,
            "stratum": None,
            "psu": None,
            "ssu": None,
            "wgt": None,
        }

        def _process_component(spec: str | list[str] | tuple[str, ...] | None, name: str):
            if not spec:
                return None, None
            if isinstance(spec, str):
                target_col = spec
            elif isinstance(spec, (list, tuple)) and len(spec) == 1:
                target_col = spec[0]
            else:
                cols = list(spec)
                target_col = f"{name}{_INTERNAL_CONCAT_SUFFIX}"
                if target_col not in local_data.columns:
                    expr = pl.concat_str(
                        [pl.col(c).cast(pl.Utf8).fill_null("__Null__") for c in cols],
                        separator=_BY_SEP,
                    )
                    s_temp = local_data.select(expr.alias(target_col))[target_col]
                else:
                    s_temp = local_data[target_col]

                if s_temp.dtype in (
                    pl.Int8,
                    pl.Int16,
                    pl.Int32,
                    pl.Int64,
                    pl.UInt8,
                    pl.UInt32,
                    pl.UInt64,
                ):
                    return s_temp.to_numpy(), None
                return (
                    s_temp.cast(pl.Categorical).to_physical().to_numpy(),
                    s_temp.unique().to_list(),
                )

            s = local_data[target_col]
            if s.dtype in (pl.Int8, pl.Int16, pl.Int32, pl.Int64, pl.UInt8, pl.UInt32, pl.UInt64):
                return s.to_numpy(), None
            return (s.cast(pl.Categorical).to_physical().to_numpy(), s.unique().to_list())

        cache["stratum"] = _process_component(design.stratum, "stratum")
        cache["psu"] = _process_component(design.psu, "psu")

        if design.ssu:
            ssu_cols = _colspec_to_list(design.ssu)
            if len(ssu_cols) == 1:
                cache["ssu"] = local_data[ssu_cols[0]].to_numpy()
            else:
                cache["ssu"] = None

        cache["wgt"] = (
            local_data[design.wgt].to_numpy()
            if design.wgt
            else np.ones(local_data.height, dtype=float)
        )

        self._design_cache = cache
        return cache

    def _get_polars_design_info(self) -> dict[str, Any]:
        if self._polars_cache is not None:
            _d2 = self._sample._data
            _h2 = (
                cast(pl.DataFrame, _d2.collect()).height
                if isinstance(_d2, pl.LazyFrame)
                else cast(pl.DataFrame, _d2).height
            )
            if self._polars_cache["n_rows"] == _h2:
                return self._polars_cache
            else:
                self._polars_cache = None

        design = self._sample._design
        _data_raw = self._sample._data
        data: pl.DataFrame = (
            cast(pl.DataFrame, _data_raw.collect())
            if isinstance(_data_raw, pl.LazyFrame)
            else cast(pl.DataFrame, _data_raw)
        )

        singleton_result = getattr(self._sample.singleton, "last_result", None)
        config = singleton_result.config if singleton_result else None

        strata_col = None
        psu_col = None
        weight_col = design.wgt

        if config and config.var_stratum_col:
            strata_col = config.var_stratum_col
            psu_col = config.var_psu_col
        else:
            if design.stratum:
                if isinstance(design.stratum, str):
                    strata_col = design.stratum
                elif isinstance(design.stratum, (list, tuple)):
                    if len(design.stratum) == 1:
                        strata_col = design.stratum[0]
                    else:
                        strata_col = f"_strata_{_INTERNAL_CONCAT_SUFFIX}"
                        if strata_col not in data.columns:
                            expr = pl.concat_str(
                                [
                                    pl.col(c).cast(pl.String).fill_null("__Null__")
                                    for c in design.stratum
                                ],
                                separator=_BY_SEP,
                            )
                            data = data.with_columns(expr.alias(strata_col))

            if design.psu:
                if isinstance(design.psu, str):
                    psu_col = design.psu
                elif isinstance(design.psu, (list, tuple)):
                    if len(design.psu) == 1:
                        psu_col = design.psu[0]
                    else:
                        psu_col = f"_psu_{_INTERNAL_CONCAT_SUFFIX}"
                        if psu_col not in data.columns:
                            expr = pl.concat_str(
                                [
                                    pl.col(c).cast(pl.String).fill_null("__Null__")
                                    for c in design.psu
                                ],
                                separator=_BY_SEP,
                            )
                            data = data.with_columns(expr.alias(psu_col))

        casts = []
        if strata_col and data[strata_col].dtype != pl.String:
            casts.append(pl.col(strata_col).cast(pl.String))
        if psu_col and data[psu_col].dtype != pl.String:
            casts.append(pl.col(psu_col).cast(pl.String))
        if casts:
            data = data.with_columns(casts)

        if not weight_col:
            weight_col = "__unit_wgt__"
            if weight_col not in data.columns:
                data = data.with_columns(pl.lit(1.0).alias(weight_col))
        else:
            if data[weight_col].dtype != pl.Float64:
                data = data.with_columns(pl.col(weight_col).cast(pl.Float64))

        _fin = self._sample._data
        _fin_h = (
            cast(pl.DataFrame, _fin.collect()).height
            if isinstance(_fin, pl.LazyFrame)
            else cast(pl.DataFrame, _fin).height
        )
        self._polars_cache = {
            "n_rows": _fin_h,
            "data": data,
            "strata_col": strata_col,
            "psu_col": psu_col,
            "weight_col": weight_col,
            "singleton_config": config,
        }
        return self._polars_cache

    def _ensure_float64(self, data: pl.DataFrame, cols: list[str]) -> pl.DataFrame:
        casts = [
            pl.col(c).cast(pl.Float64)
            for c in cols
            if c in data.columns and data[c].dtype != pl.Float64
        ]
        return data.with_columns(casts) if casts else data

    def _get_enum_value(self, config, attr="method") -> str:
        if not hasattr(config, attr):
            return ""
        val = getattr(config, attr)
        return str(val.value) if hasattr(val, "value") else str(val)

    def _prepare_polars_data(
        self,
        y: str,
        x: str | None = None,
        by: str | Sequence[str] | None = None,
        drop_nulls: bool = False,
        cast_y_float: bool = True,
        apply_singleton_filter: bool = True,
    ) -> tuple[pl.DataFrame, str | None, str | None, str | None, str | None, list[str]]:
        cache = self._get_polars_design_info()
        data = cache["data"]
        strata_col = cache["strata_col"]
        psu_col = cache["psu_col"]
        weight_col = cache["weight_col"]
        singleton_config = cache.get("singleton_config")

        by_cols = _colspec_to_list(by)
        if len(by_cols) > 1:
            by_col = f"_by_{_INTERNAL_CONCAT_SUFFIX}"
            if by_col not in data.columns:
                expr = pl.concat_str(
                    [pl.col(c).cast(pl.String).fill_null("__Null__") for c in by_cols],
                    separator=_BY_SEP,
                )
                data = data.with_columns(expr.alias(by_col))
        elif len(by_cols) == 1:
            by_col = by_cols[0]
        else:
            by_col = None

        filter_cols = [y]
        if x:
            filter_cols.append(x)
        if weight_col and weight_col != "__unit_wgt__":
            filter_cols.append(weight_col)
        if by_col:
            filter_cols.append(by_col)

        if drop_nulls:
            data = data.filter(pl.all_horizontal(pl.col(c).is_not_null() for c in filter_cols))

        if apply_singleton_filter and singleton_config and singleton_config.var_exclude_col:
            if singleton_config.var_exclude_col in data.columns:
                data = data.filter(~pl.col(singleton_config.var_exclude_col))

        casts = []
        if cast_y_float:
            casts.append(pl.col(y).cast(pl.Float64))
        if x:
            casts.append(pl.col(x).cast(pl.Float64))
        if by_col and data[by_col].dtype != pl.String:
            casts.append(pl.col(by_col).cast(pl.String))
        if casts:
            data = data.with_columns(casts)

        return data, strata_col, psu_col, weight_col, by_col, by_cols

    def _adjust_variance_for_singletons(
        self, result_df: pl.DataFrame, param: PopParam = PopParam.TOTAL
    ) -> pl.DataFrame:
        """Legacy helper for tests. New code uses _apply_scale_adjustment."""
        return self._apply_scale_adjustment(result_df, result_df, param=param)

    def _get_center_method(self) -> str | None:
        cache = self._get_polars_design_info()
        config = cache.get("singleton_config")
        if config:
            method_str = self._get_enum_value(config, "method").lower()
            if method_str in ("center", "adjust"):
                return "center"
        return None

    def _should_run_double_pass(self) -> bool:
        cache = self._get_polars_design_info()
        config = cache.get("singleton_config")
        if config:
            return self._get_enum_value(config, "method").lower() == "scale"
        return False

    def _apply_scale_adjustment(
        self, full_df: pl.DataFrame, filtered_df: pl.DataFrame, param: PopParam = PopParam.TOTAL
    ) -> pl.DataFrame:
        """
        Merge Point Estimate (from full_df) and Variance (from filtered_df)
        and apply inflation factor for SCALE method.
        """
        cache = self._get_polars_design_info()
        config = cache.get("singleton_config")

        if not config:
            return filtered_df

        method_str = self._get_enum_value(config, "method").lower()
        if method_str != "scale":
            return filtered_df

        f = config.singleton_fraction
        if f is None or f >= 1.0:
            return filtered_df

        if param == PopParam.TOTAL:
            inflation_factor = 1.0 / (1.0 - f)
        else:
            inflation_factor = 1.0 - f

        sqrt_factor = math.sqrt(inflation_factor)

        if full_df is filtered_df:
            merged = filtered_df
        else:
            std_cols = {"y", "est", "se", "var", "df", "n", "deff", "level"}
            extra_cols = [c for c in filtered_df.columns if c not in std_cols]
            by_col_name = extra_cols[0] if extra_cols else None

            join_on = ["y"]
            if by_col_name:
                join_on.append(by_col_name)
            if "level" in filtered_df.columns:
                join_on.append("level")

            full_subset = full_df.select(join_on + ["est"])

            if "est" in filtered_df.columns:
                merged = filtered_df.drop("est").join(full_subset, on=join_on, how="left")
            else:
                merged = filtered_df.join(full_subset, on=join_on, how="left")

        return merged.with_columns(
            (pl.col("var") * inflation_factor).alias("var"),
            (pl.col("se") * sqrt_factor).alias("se"),
        )

    def _compute_prop_ci(
        self,
        p: float,
        se: float,
        t_crit: float,
        method: str,
    ) -> tuple[float, float]:
        """Compute confidence interval for a proportion."""
        if method == "logit":
            if p <= 0 or p >= 1:
                return (p, p)
            scale = se / (p * (1.0 - p)) if se > 0 else 0
            logit = math.log(p / (1 - p))
            lci = 1.0 / (1.0 + math.exp(-(logit - t_crit * scale)))
            uci = 1.0 / (1.0 + math.exp(-(logit + t_crit * scale)))
            return (lci, uci)

        elif method == "beta":
            from scipy.stats import beta as beta_dist
            from scipy.stats import norm

            if p <= 0 or p >= 1 or se <= 0:
                return (p, p)

            n_eff = (p * (1 - p)) / (se**2)
            alpha_param = max(0.5, p * n_eff)
            beta_param = max(0.5, (1 - p) * n_eff)
            prob_lower = norm.cdf(-t_crit)
            prob_upper = norm.cdf(t_crit)
            lci = beta_dist.ppf(prob_lower, alpha_param, beta_param)
            uci = beta_dist.ppf(prob_upper, alpha_param, beta_param)
            return (lci, uci)

        else:
            raise ValueError(f"Unknown CI method: {method}")

    def _polars_result_to_param_est(
        self,
        result_df: pl.DataFrame,
        y_name: str,
        param: PopParam,
        alpha: float,
        deff: bool,
        by_col: str | None,
        as_factor: bool,
        x_name: str | None = None,
        ci_method: str = "logit",
    ) -> list[ParamEst]:
        est_list = []

        for row in result_df.iter_rows(named=True):
            est = row["est"]
            se = row["se"]
            df_val = row["df"]

            cv = se / est if est != 0 else float("inf")
            t_crit = stats.t.ppf(1 - alpha / 2, df_val) if df_val > 0 else 1.96

            is_prop = (param == PopParam.PROP) or as_factor
            if is_prop:
                lci, uci = self._compute_prop_ci(p=est, se=se, t_crit=t_crit, method=ci_method)
            else:
                lci = est - t_crit * se
                uci = est + t_crit * se

            deff_val = row.get("deff") if deff else None

            by_level = None
            if by_col and by_col in row:
                by_level = (row[by_col],)

            y_level = None
            if as_factor and "level" in row:
                level_val = row["level"]
                try:
                    y_level = int(level_val)
                except (ValueError, TypeError):
                    y_level = level_val

            p = ParamEst(
                y=y_name,
                est=est,
                se=se,
                cv=cv,
                lci=lci,
                uci=uci,
                deff=deff_val,
                by=(by_col,) if by_col else None,
                by_level=by_level,
                y_level=y_level,
                x=x_name,
            )
            est_list.append(p)

        return est_list

    def _median_result_to_param_est(
        self,
        result_df: pl.DataFrame,
        y_name: str,
        alpha: float,
        by_col: str | None,
        data: pl.DataFrame,
        weight_col: str,
    ) -> list[ParamEst]:
        """
        Convert Rust median result to ParamEst list.
        Uses Woodruff method to compute CI by inverting the CDF.
        """
        est_list = []

        for row in result_df.iter_rows(named=True):
            est = row["est"]
            se_p = row["se"]
            df_val = row["df"]

            t_crit = stats.t.ppf(1 - alpha / 2, df_val) if df_val > 0 else 1.96

            # Woodruff CI: invert CDF at p ± t*se_p
            p_lower = max(0.0, 0.5 - t_crit * se_p)
            p_upper = min(1.0, 0.5 + t_crit * se_p)

            # Get domain filter if applicable
            if by_col and by_col in row:
                domain_val = row[by_col]
                domain_data = data.filter(pl.col(by_col) == domain_val)
            else:
                domain_data = data

            # Invert CDF to get quantile CI
            lci = self._invert_cdf(domain_data, y_name, weight_col, p_lower)
            uci = self._invert_cdf(domain_data, y_name, weight_col, p_upper)

            # SE on quantile scale (approximate)
            se_q = (uci - lci) / (2 * t_crit) if t_crit > 0 else se_p
            cv = se_q / est if est != 0 else float("inf")

            by_level = None
            if by_col and by_col in row:
                by_level = (row[by_col],)

            p = ParamEst(
                y=y_name,
                est=est,
                se=se_q,
                cv=cv,
                lci=lci,
                uci=uci,
                deff=None,
                by=(by_col,) if by_col else None,
                by_level=by_level,
                y_level=None,
                x=None,
            )
            est_list.append(p)

        return est_list

    def _invert_cdf(
        self,
        data: pl.DataFrame,
        y_col: str,
        weight_col: str,
        p: float,
    ) -> float:
        """Invert the weighted CDF to find the quantile at probability p."""
        df = data.select([y_col, weight_col]).drop_nulls()
        if df.height == 0:
            return float("nan")

        df = df.sort(y_col)
        y_vals = df[y_col].to_numpy()
        w_vals = df[weight_col].to_numpy()

        cumsum = np.cumsum(w_vals)
        total = cumsum[-1]
        if total <= 0:
            return float("nan")

        cdf = cumsum / total

        if p <= 0:
            return float(y_vals[0])
        if p >= 1:
            return float(y_vals[-1])

        idx = np.searchsorted(cdf, p)

        if idx == 0:
            return float(y_vals[0])
        if idx >= len(y_vals):
            return float(y_vals[-1])

        p_low = cdf[idx - 1]
        p_high = cdf[idx]
        y_low = y_vals[idx - 1]
        y_high = y_vals[idx]

        if p_high == p_low:
            return float(y_high)

        frac = (p - p_low) / (p_high - p_low)
        return float(y_low + frac * (y_high - y_low))

    def _replicate_median_result_to_param_est(
        self,
        result_df: pl.DataFrame,
        y_name: str,
        alpha: float,
        by_col: str | None,
    ) -> list[ParamEst]:
        """Convert Rust replicate median result to ParamEst list."""
        est_list = []

        for row in result_df.iter_rows(named=True):
            est = row["est"]
            se = row["se"]
            df_val = row["df"]

            cv = se / est if est != 0 else float("inf")
            t_crit = stats.t.ppf(1 - alpha / 2, df_val) if df_val > 0 else 1.96

            lci = est - t_crit * se
            uci = est + t_crit * se

            by_level = None
            if by_col and by_col in row:
                by_level = (row[by_col],)

            p = ParamEst(
                y=y_name,
                est=est,
                se=se,
                cv=cv,
                lci=lci,
                uci=uci,
                deff=None,
                by=(by_col,) if by_col else None,
                by_level=by_level,
                y_level=None,
                x=None,
            )
            est_list.append(p)

        return est_list

    def _build_estimate_result_light(
        self,
        est_list: list[ParamEst],
        est_cov: np.ndarray,
        param: PopParam,
        alpha: float,
        by_cols: list[str],
        as_factor: bool,
        method: EstimationMethod = EstimationMethod.TAYLOR,
    ) -> Estimate:
        metadata = getattr(self._sample, "_metadata", None)
        estimate = Estimate(param, alpha=alpha, metadata=metadata)
        estimate.method = method
        estimate.covariance = est_cov
        estimate.as_factor = as_factor

        if by_cols and len(by_cols) > 0:
            by_tuple = tuple(by_cols)
            final_ests = []
            for p in est_list:
                new_p = ParamEst(
                    y=p.y,
                    est=p.est,
                    se=p.se,
                    cv=p.cv,
                    lci=p.lci,
                    uci=p.uci,
                    deff=p.deff,
                    by=by_tuple,
                    by_level=p.by_level,
                    y_level=p.y_level,
                    x=p.x,
                    x_level=p.x_level,
                )
                final_ests.append(new_p)
            estimate.estimates = final_ests
        else:
            estimate.estimates = est_list

        d_cache = self._get_factorized_design()

        if d_cache["stratum"] is not None:
            strata_info = d_cache["stratum"]
            if isinstance(strata_info, tuple) and strata_info[1] is not None:
                estimate.strata = strata_info[1]
                estimate.n_strata = len(strata_info[1])
            else:
                arr = strata_info[0] if isinstance(strata_info, tuple) else strata_info
                estimate.strata = np.unique(arr).tolist()
                estimate.n_strata = len(estimate.strata)

        if d_cache["psu"] is not None:
            psu_info = d_cache["psu"]
            if isinstance(psu_info, tuple) and psu_info[1] is not None:
                estimate.n_psus = len(psu_info[1])
            else:
                arr = psu_info[0] if isinstance(psu_info, tuple) else psu_info
                estimate.n_psus = len(np.unique(arr))
        elif d_cache["wgt"] is not None:
            estimate.n_psus = len(d_cache["wgt"])

        estimate.degrees_freedom = max(0, (estimate.n_psus or 0) - (estimate.n_strata or 0))

        return estimate

    def _resolve_method(self, user_method: EstimationMethod | None) -> EstimationMethod:
        if user_method is not None:
            return user_method
        design_rw = self._sample._design.rep_wgts
        if design_rw is not None:
            return design_rw.method
        return EstimationMethod.TAYLOR

    # ----------------------------------------------------------------
    # Domain Estimation Helpers
    # ----------------------------------------------------------------
    def _where_to_string(self, where: WhereArg) -> str:
        """Convert where clause to a human-readable string for display."""
        if where is None:
            return ""

        def _clean_expr_str(expr: pl.Expr | ExprLike) -> str:
            if hasattr(expr, "_e"):
                expr = expr._e

            s = str(expr)
            s = s.strip()
            if s.startswith("[") and s.endswith("]"):
                s = s[1:-1].strip()

            s = re.sub(r'\(col\("([^"]+)"\)\)', r"\1", s)
            s = re.sub(r'col\("([^"]+)"\)', r"\1", s)
            s = re.sub(r"\(dyn int: (-?\d+)\)", r"\1", s)
            s = re.sub(r"\(dyn float: (-?[\d.]+(?:e[+-]?\d+)?)\)", r"\1", s)
            s = re.sub(r'\((?:String|Utf8): ("[^"]*")\)', r"\1", s)
            s = re.sub(r"\(Boolean: (true|false)\)", r"\1", s)
            s = re.sub(r"\.strict_cast\([^)]*\)", "", s)
            s = re.sub(r"\s+", " ", s).strip()

            return s

        if hasattr(where, "_e"):
            return _clean_expr_str(where)

        if isinstance(where, pl.Expr):
            return _clean_expr_str(where)

        if isinstance(where, Sequence) and not isinstance(where, (str, Mapping)):
            parts = [_clean_expr_str(w) for w in where]
            return " & ".join(parts)

        if isinstance(where, Mapping):
            parts = []
            for k, v in where.items():
                if isinstance(v, (list, tuple, set)):
                    parts.append(f"{k} in {list(v)!r}")
                else:
                    parts.append(f"{k} == {v!r}")
            return " & ".join(parts)

        return str(where)

    def _compile_where_expr(self, where: WhereArg) -> pl.Expr:
        """Compile where clause to Polars expression using wrangling helper."""
        result = self._sample.wrangling._compile_where_to_pl_expr(where)
        return cast(pl.Expr, result)

    def _with_domain_estimation(
        self,
        where: WhereArg,
        by: str | Sequence[str] | None,
        estimator: Callable[[str | Sequence[str] | None], Estimate],
    ) -> Estimate:
        """
        Execute estimation with domain restriction (subpopulation analysis).
        """
        pred = self._compile_where_expr(where)
        _orig_raw = self._sample._data
        original_data: pl.DataFrame = (
            cast(pl.DataFrame, _orig_raw.collect())
            if isinstance(_orig_raw, pl.LazyFrame)
            else cast(pl.DataFrame, _orig_raw)
        )

        self._sample._data = original_data.with_columns(
            pl.when(pred).then(pl.lit("1")).otherwise(pl.lit("0")).alias(_DOMAIN_COL)
        )

        original_by_cols = _colspec_to_list(by)

        if by is None:
            combined_by = _DOMAIN_COL
        elif isinstance(by, str):
            combined_by = [_DOMAIN_COL, by]
        else:
            combined_by = [_DOMAIN_COL, *list(by)]

        try:
            result = estimator(combined_by)
            result = self._filter_domain_estimates(result, original_by_cols)
            result.where_clause = self._where_to_string(where)
            return result

        finally:
            self._sample._data = original_data
            self._design_cache = None
            self._polars_cache = None

    def _filter_domain_estimates(
        self,
        estimate: Estimate,
        original_by_cols: list[str],
    ) -> Estimate:
        from svy.core.constants import _BY_SEP

        original_by_tuple = tuple(original_by_cols) if original_by_cols else None
        has_original_by = len(original_by_cols) > 0

        filtered_estimates = []
        for p in estimate.estimates:
            if p.by_level is None or len(p.by_level) == 0:
                continue

            by_level_str = str(p.by_level[0])

            if has_original_by:
                parts = by_level_str.split(_BY_SEP, maxsplit=len(original_by_cols))
                if len(parts) < 2:
                    continue
                domain_val = parts[0]
                by_parts = parts[1:]
            else:
                domain_val = by_level_str
                by_parts = []

            if domain_val != "1":
                continue

            new_by_level: tuple[str, ...] | None = tuple(by_parts) if by_parts else None

            new_p = ParamEst(
                y=p.y,
                est=p.est,
                se=p.se,
                cv=p.cv,
                lci=p.lci,
                uci=p.uci,
                deff=p.deff,
                by=original_by_tuple,
                by_level=new_by_level,
                y_level=p.y_level,
                x=p.x,
                x_level=p.x_level,
            )
            filtered_estimates.append(new_p)

        estimate.estimates = filtered_estimates
        return estimate

    # ----------------------------------------------------------------
    # Unified Public APIs
    # ----------------------------------------------------------------

    def mean(
        self,
        y: str,
        *,
        by: str | Sequence[str] | None = None,
        where: WhereArg = None,
        method: EstimationMethod | None = None,
        deff: bool = False,
        fay_coef: float = 0.0,
        as_factor: bool = False,
        variance_center: Literal["rep_mean", "estimate"] = "rep_mean",
        alpha: float = 0.05,
        drop_nulls: bool = False,
    ) -> Estimate:
        """Estimate population mean with standard errors."""
        if where is not None:
            return self._with_domain_estimation(
                where=where,
                by=by,
                estimator=lambda combined_by: self.mean(
                    y=y,
                    by=combined_by,
                    where=None,
                    method=method,
                    deff=deff,
                    fay_coef=fay_coef,
                    as_factor=as_factor,
                    variance_center=variance_center,
                    alpha=alpha,
                    drop_nulls=drop_nulls,
                ),
            )

        target_method = self._resolve_method(method)

        if target_method == EstimationMethod.TAYLOR:
            return self._taylor_mean_polars(
                y=y,
                by=by,
                deff=deff,
                alpha=alpha,
                drop_nulls=drop_nulls,
                as_factor=as_factor,
                param=PopParam.MEAN,
            )

        return self._replicate_estimate(
            method=target_method,
            param=PopParam.MEAN,
            y=y,
            by=by,
            fay_coef=fay_coef,
            as_factor=as_factor,
            variance_center=variance_center,
            alpha=alpha,
            drop_nulls=drop_nulls,
        )

    def total(
        self,
        y: str,
        *,
        by: str | Sequence[str] | None = None,
        where: WhereArg = None,
        method: EstimationMethod | None = None,
        deff: bool = False,
        fay_coef: float = 0.0,
        as_factor: bool = False,
        variance_center: Literal["rep_mean", "estimate"] = "rep_mean",
        alpha: float = 0.05,
        drop_nulls: bool = False,
    ) -> Estimate:
        """Estimate population total with standard errors."""
        if where is not None:
            return self._with_domain_estimation(
                where=where,
                by=by,
                estimator=lambda combined_by: self.total(
                    y=y,
                    by=combined_by,
                    where=None,
                    method=method,
                    deff=deff,
                    fay_coef=fay_coef,
                    as_factor=as_factor,
                    variance_center=variance_center,
                    alpha=alpha,
                    drop_nulls=drop_nulls,
                ),
            )

        target_method = self._resolve_method(method)

        if target_method == EstimationMethod.TAYLOR:
            return self._taylor_total_polars(
                y=y,
                by=by,
                deff=deff,
                alpha=alpha,
                drop_nulls=drop_nulls,
                as_factor=as_factor,
            )

        return self._replicate_estimate(
            method=target_method,
            param=PopParam.TOTAL,
            y=y,
            by=by,
            fay_coef=fay_coef,
            as_factor=as_factor,
            variance_center=variance_center,
            alpha=alpha,
            drop_nulls=drop_nulls,
        )

    def prop(
        self,
        y: str,
        *,
        by: str | Sequence[str] | None = None,
        where: WhereArg = None,
        method: EstimationMethod | None = None,
        ci_method: Literal["logit", "beta"] = "logit",
        deff: bool = False,
        fay_coef: float = 0.0,
        variance_center: Literal["rep_mean", "estimate"] = "rep_mean",
        alpha: float = 0.05,
        drop_nulls: bool = False,
    ) -> Estimate:
        """Estimate population proportion with standard errors."""
        if where is not None:
            return self._with_domain_estimation(
                where=where,
                by=by,
                estimator=lambda combined_by: self.prop(
                    y=y,
                    by=combined_by,
                    where=None,
                    method=method,
                    ci_method=ci_method,
                    deff=deff,
                    fay_coef=fay_coef,
                    variance_center=variance_center,
                    alpha=alpha,
                    drop_nulls=drop_nulls,
                ),
            )

        target_method = self._resolve_method(method)

        if target_method == EstimationMethod.TAYLOR:
            return self._taylor_prop_polars(
                y=y,
                by=by,
                deff=deff,
                alpha=alpha,
                ci_method=ci_method,
                drop_nulls=drop_nulls,
            )

        return self._replicate_estimate(
            method=target_method,
            param=PopParam.PROP,
            y=y,
            by=by,
            fay_coef=fay_coef,
            as_factor=True,
            variance_center=variance_center,
            alpha=alpha,
            ci_method=ci_method,
            drop_nulls=drop_nulls,
        )

    def ratio(
        self,
        y: str,
        x: str,
        *,
        by: str | Sequence[str] | None = None,
        where: WhereArg = None,
        method: EstimationMethod | None = None,
        deff: bool = False,
        fay_coef: float = 0.0,
        variance_center: Literal["rep_mean", "estimate"] = "rep_mean",
        alpha: float = 0.05,
        drop_nulls: bool = False,
    ) -> Estimate:
        """Estimate population ratio (y/x) with standard errors."""
        if where is not None:
            return self._with_domain_estimation(
                where=where,
                by=by,
                estimator=lambda combined_by: self.ratio(
                    y=y,
                    x=x,
                    by=combined_by,
                    where=None,
                    method=method,
                    deff=deff,
                    fay_coef=fay_coef,
                    variance_center=variance_center,
                    alpha=alpha,
                    drop_nulls=drop_nulls,
                ),
            )

        target_method = self._resolve_method(method)

        if target_method == EstimationMethod.TAYLOR:
            return self._taylor_ratio_polars(
                y=y,
                x=x,
                by=by,
                deff=deff,
                alpha=alpha,
                drop_nulls=drop_nulls,
            )

        return self._replicate_estimate(
            method=target_method,
            param=PopParam.RATIO,
            y=y,
            x=x,
            by=by,
            fay_coef=fay_coef,
            variance_center=variance_center,
            alpha=alpha,
            drop_nulls=drop_nulls,
        )

    def median(
        self,
        y: str,
        *,
        by: str | Sequence[str] | None = None,
        where: WhereArg = None,
        method: EstimationMethod | None = None,
        fay_coef: float = 0.0,
        q_method: QuantileMethod = QuantileMethod.HIGHER,
        variance_center: Literal["rep_mean", "estimate"] = "rep_mean",
        alpha: float = 0.05,
        drop_nulls: bool = False,
    ) -> Estimate:
        """Estimate population median with standard errors."""
        if where is not None:
            return self._with_domain_estimation(
                where=where,
                by=by,
                estimator=lambda combined_by: self.median(
                    y=y,
                    by=combined_by,
                    where=None,
                    method=method,
                    fay_coef=fay_coef,
                    q_method=q_method,
                    variance_center=variance_center,
                    alpha=alpha,
                    drop_nulls=drop_nulls,
                ),
            )

        target_method = self._resolve_method(method)

        if target_method == EstimationMethod.TAYLOR:
            return self._taylor_median_polars(
                y=y,
                by=by,
                q_method=q_method,
                alpha=alpha,
                drop_nulls=drop_nulls,
            )

        return self._replicate_median_polars(
            y=y,
            method=target_method,
            by=by,
            fay_coef=fay_coef,
            q_method=q_method,
            variance_center=variance_center,
            alpha=alpha,
            drop_nulls=drop_nulls,
        )

    # ----------------------------------------------------------------
    # Internal Taylor Helpers - Polars Backend
    # ----------------------------------------------------------------

    def _taylor_mean_polars(
        self,
        y: str,
        *,
        param: PopParam = PopParam.MEAN,
        by: str | Sequence[str] | None = None,
        deff: bool = False,
        as_factor: bool = False,
        alpha: float = 0.05,
        drop_nulls: bool = False,
    ) -> Estimate:
        data, strata_col, psu_col, weight_col, by_col, by_cols = self._prepare_polars_data(
            y=y, by=by, drop_nulls=drop_nulls, cast_y_float=True, apply_singleton_filter=True
        )

        center_arg = self._get_center_method()
        fn = rs.taylor_prop if as_factor else rs.taylor_mean

        result_df = fn(
            data,
            value_col=y,
            weight_col=weight_col,
            strata_col=strata_col,
            psu_col=psu_col,
            by_col=by_col,
            singleton_method=center_arg,
        )

        if self._should_run_double_pass():
            data_full, _, _, _, _, _ = self._prepare_polars_data(
                y=y, by=by, drop_nulls=drop_nulls, cast_y_float=True, apply_singleton_filter=False
            )
            result_full = fn(
                data_full,
                value_col=y,
                weight_col=weight_col,
                strata_col=strata_col,
                psu_col=psu_col,
                by_col=by_col,
                singleton_method=center_arg,
            )
            result_df = self._apply_scale_adjustment(result_full, result_df, param=param)

        est_list = self._polars_result_to_param_est(
            result_df, y, param, alpha, deff, by_col, as_factor
        )
        est_cov = np.diag(result_df["var"].to_numpy())
        return self._build_estimate_result_light(
            est_list, est_cov, param, alpha, by_cols, as_factor, method=EstimationMethod.TAYLOR
        )

    def _taylor_total_polars(
        self,
        y: str,
        *,
        by: str | Sequence[str] | None = None,
        deff: bool = False,
        as_factor: bool = False,
        alpha: float = 0.05,
        drop_nulls: bool = False,
    ) -> Estimate:
        data, strata_col, psu_col, weight_col, by_col, by_cols = self._prepare_polars_data(
            y=y, by=by, drop_nulls=drop_nulls, cast_y_float=True, apply_singleton_filter=True
        )
        center_arg = self._get_center_method()

        result_df = rs.taylor_total(
            data,
            value_col=y,
            weight_col=weight_col,
            strata_col=strata_col,
            psu_col=psu_col,
            by_col=by_col,
            singleton_method=center_arg,
        )

        if self._should_run_double_pass():
            data_full, _, _, _, _, _ = self._prepare_polars_data(
                y=y, by=by, drop_nulls=drop_nulls, cast_y_float=True, apply_singleton_filter=False
            )
            result_full = rs.taylor_total(
                data_full,
                value_col=y,
                weight_col=weight_col,
                strata_col=strata_col,
                psu_col=psu_col,
                by_col=by_col,
                singleton_method=center_arg,
            )
            result_df = self._apply_scale_adjustment(result_full, result_df, param=PopParam.TOTAL)

        est_list = self._polars_result_to_param_est(
            result_df, y, PopParam.TOTAL, alpha, deff, by_col, as_factor=False
        )
        est_cov = np.diag(result_df["var"].to_numpy())
        return self._build_estimate_result_light(
            est_list,
            est_cov,
            PopParam.TOTAL,
            alpha,
            by_cols,
            as_factor=False,
            method=EstimationMethod.TAYLOR,
        )

    def _taylor_ratio_polars(
        self,
        y: str,
        x: str,
        *,
        by: str | Sequence[str] | None = None,
        deff: bool = False,
        alpha: float = 0.05,
        drop_nulls: bool = False,
    ) -> Estimate:
        data, strata_col, psu_col, weight_col, by_col, by_cols = self._prepare_polars_data(
            y=y, x=x, by=by, drop_nulls=drop_nulls, cast_y_float=True, apply_singleton_filter=True
        )
        center_arg = self._get_center_method()

        result_df = rs.taylor_ratio(
            data,
            numerator_col=y,
            denominator_col=x,
            weight_col=weight_col,
            strata_col=strata_col,
            psu_col=psu_col,
            by_col=by_col,
            singleton_method=center_arg,
        )

        if self._should_run_double_pass():
            data_full, _, _, _, _, _ = self._prepare_polars_data(
                y=y,
                x=x,
                by=by,
                drop_nulls=drop_nulls,
                cast_y_float=True,
                apply_singleton_filter=False,
            )
            result_full = rs.taylor_ratio(
                data_full,
                numerator_col=y,
                denominator_col=x,
                weight_col=weight_col,
                strata_col=strata_col,
                psu_col=psu_col,
                by_col=by_col,
                singleton_method=center_arg,
            )
            result_df = self._apply_scale_adjustment(result_full, result_df, param=PopParam.RATIO)

        est_list = self._polars_result_to_param_est(
            result_df, y, PopParam.RATIO, alpha, deff, by_col, as_factor=False, x_name=x
        )
        est_cov = np.diag(result_df["var"].to_numpy())
        return self._build_estimate_result_light(
            est_list,
            est_cov,
            PopParam.RATIO,
            alpha,
            by_cols,
            as_factor=False,
            method=EstimationMethod.TAYLOR,
        )

    def _taylor_prop_polars(
        self,
        y: str,
        *,
        by: str | Sequence[str] | None = None,
        deff: bool = False,
        alpha: float = 0.05,
        ci_method: str = "logit",
        drop_nulls: bool = False,
    ) -> Estimate:
        data, strata_col, psu_col, weight_col, by_col, by_cols = self._prepare_polars_data(
            y=y, by=by, drop_nulls=drop_nulls, cast_y_float=False, apply_singleton_filter=True
        )
        center_arg = self._get_center_method()

        result_df = rs.taylor_prop(
            data,
            value_col=y,
            weight_col=weight_col,
            strata_col=strata_col,
            psu_col=psu_col,
            by_col=by_col,
            singleton_method=center_arg,
        )

        if self._should_run_double_pass():
            data_full, _, _, _, _, _ = self._prepare_polars_data(
                y=y, by=by, drop_nulls=drop_nulls, cast_y_float=False, apply_singleton_filter=False
            )
            result_full = rs.taylor_prop(
                data_full,
                value_col=y,
                weight_col=weight_col,
                strata_col=strata_col,
                psu_col=psu_col,
                by_col=by_col,
                singleton_method=center_arg,
            )
            result_df = self._apply_scale_adjustment(result_full, result_df, param=PopParam.PROP)

        est_list = self._polars_result_to_param_est(
            result_df, y, PopParam.PROP, alpha, deff, by_col, as_factor=True, ci_method=ci_method
        )
        est_cov = np.diag(result_df["var"].to_numpy())
        return self._build_estimate_result_light(
            est_list,
            est_cov,
            PopParam.PROP,
            alpha,
            by_cols,
            as_factor=True,
            method=EstimationMethod.TAYLOR,
        )

    def _taylor_median_polars(
        self,
        y: str,
        *,
        by: str | Sequence[str] | None = None,
        q_method: QuantileMethod = QuantileMethod.HIGHER,
        alpha: float = 0.05,
        drop_nulls: bool = False,
    ) -> Estimate:
        """Taylor linearization median using Rust backend."""
        data, strata_col, psu_col, weight_col, by_col, by_cols = self._prepare_polars_data(
            y=y, by=by, drop_nulls=drop_nulls, cast_y_float=True, apply_singleton_filter=True
        )

        if not drop_nulls and data[y].null_count() > 0:
            raise ValueError(f"Missing values found in '{y}'. Use drop_nulls=True to ignore.")

        center_arg = self._get_center_method()

        q_method_str = q_method.value if hasattr(q_method, "value") else str(q_method).lower()

        result_df = rs.taylor_median(
            data,
            value_col=y,
            weight_col=weight_col,
            strata_col=strata_col,
            psu_col=psu_col,
            by_col=by_col,
            singleton_method=center_arg,
            quantile_method=q_method_str,
        )

        if weight_col is None:
            raise ValueError("weight_col is required for median estimation")
        est_list = self._median_result_to_param_est(result_df, y, alpha, by_col, data, weight_col)
        est_cov = np.diag(result_df["var"].to_numpy())

        return self._build_estimate_result_light(
            est_list,
            est_cov,
            PopParam.MEDIAN,
            alpha,
            by_cols,
            as_factor=False,
            method=EstimationMethod.TAYLOR,
        )

    # ----------------------------------------------------------------
    # Internal Replication Helpers
    # ----------------------------------------------------------------
    def _replicate_estimate(
        self,
        method,
        param,
        y,
        *,
        x=None,
        by=None,
        fay_coef=0.0,
        as_factor=False,
        variance_center="rep_mean",
        alpha=0.05,
        ci_method="logit",
        drop_nulls=False,
    ):
        if method not in (
            EstimationMethod.BRR,
            EstimationMethod.BOOTSTRAP,
            EstimationMethod.JACKKNIFE,
            EstimationMethod.SDR,
        ):
            raise ValueError(f"Method {method} is not a valid replication method.")
        if param == PopParam.MEAN:
            return self._replicate_mean_polars(
                y=y,
                method=method,
                by=by,
                fay_coef=fay_coef,
                variance_center=variance_center,
                alpha=alpha,
                drop_nulls=drop_nulls,
            )
        elif param == PopParam.TOTAL:
            return self._replicate_total_polars(
                y=y,
                method=method,
                by=by,
                fay_coef=fay_coef,
                variance_center=variance_center,
                alpha=alpha,
                drop_nulls=drop_nulls,
            )
        elif param == PopParam.RATIO:
            if x is None:
                raise ValueError("x must be provided for ratio estimation.")
            return self._replicate_ratio_polars(
                y=y,
                x=x,
                method=method,
                by=by,
                fay_coef=fay_coef,
                variance_center=variance_center,
                alpha=alpha,
                drop_nulls=drop_nulls,
            )
        elif param == PopParam.PROP:
            return self._replicate_prop_polars(
                y=y,
                method=method,
                by=by,
                fay_coef=fay_coef,
                variance_center=variance_center,
                alpha=alpha,
                ci_method=ci_method,
                drop_nulls=drop_nulls,
            )
        else:
            raise ValueError(f"Unsupported parameter {param} for replication estimation.")

    def _get_rep_weight_cols(self) -> list[str]:
        rw = self._sample._design.rep_wgts
        if rw is None:
            return []
        _lraw = self._sample._data
        local_data: pl.DataFrame = (
            cast(pl.DataFrame, _lraw.collect())
            if isinstance(_lraw, pl.LazyFrame)
            else cast(pl.DataFrame, _lraw)
        )
        if rw.prefix:

            def natural_keys(text):
                return [int(c) if c.isdigit() else c for c in re.split(r"(\d+)", text)]

            return sorted(
                [c for c in local_data.columns if c.startswith(rw.prefix) and c != rw.prefix],
                key=natural_keys,
            )
        elif hasattr(rw, "wgts") and rw.wgts:
            return list(cast(list[str], rw.wgts))
        return []

    def _get_rep_method_str(self, method: EstimationMethod) -> str:
        method_map = {
            EstimationMethod.BRR: "BRR",
            EstimationMethod.BOOTSTRAP: "Bootstrap",
            EstimationMethod.JACKKNIFE: "Jackknife",
            EstimationMethod.SDR: "SDR",
        }
        return method_map.get(method, "jackknife")

    def _replicate_mean_polars(
        self,
        y,
        *,
        method,
        by=None,
        fay_coef=0.0,
        variance_center="rep_mean",
        alpha=0.05,
        drop_nulls=False,
    ):
        design = self._sample._design
        _raw = self._sample._data
        data: pl.DataFrame = (
            cast(pl.DataFrame, _raw.collect())
            if isinstance(_raw, pl.LazyFrame)
            else cast(pl.DataFrame, _raw)
        )
        rw = design.rep_wgts
        if rw is None:
            raise ValueError("Replication weights required for replication-based estimation.")
        rep_weight_cols = self._get_rep_weight_cols()
        if not rep_weight_cols:
            raise ValueError("No replicate weight columns found.")
        n_reps = len(rep_weight_cols)
        df_val = int(rw.df) if rw.df and rw.df > 0 else max(1, n_reps - 1)
        final_fay = float(fay_coef) if fay_coef != 0.0 else float(rw.fay_coef)

        by_cols = _colspec_to_list(by)

        if len(by_cols) > 1:
            by_col = f"_by_{_INTERNAL_CONCAT_SUFFIX}"
            if by_col not in data.columns:
                expr = pl.concat_str(
                    [pl.col(c).cast(pl.String).fill_null("__Null__") for c in by_cols],
                    separator=_BY_SEP,
                )
                data = data.with_columns(expr.alias(by_col))
        elif len(by_cols) == 1:
            by_col = by_cols[0]
        else:
            by_col = None

        data = data.with_columns(pl.col(y).cast(pl.Float64))
        data = self._ensure_float64(data, rep_weight_cols)
        filter_cols = [y]
        weight_col = design.wgt
        if weight_col:
            filter_cols.append(weight_col)
        if by_col:
            filter_cols.append(by_col)
        if drop_nulls:
            data = data.filter(pl.all_horizontal(pl.col(c).is_not_null() for c in filter_cols))
        if by_col and data[by_col].dtype != pl.String:
            data = data.with_columns(pl.col(by_col).cast(pl.String))
        if not weight_col:
            weight_col = "__unit_wgt__"
            data = data.with_columns(pl.lit(1.0).alias(weight_col))
        else:
            data = self._ensure_float64(data, [weight_col])
        method_str = self._get_rep_method_str(method)
        result_df = rs.replicate_mean(
            data,
            value_col=y,
            weight_col=weight_col,
            rep_weight_cols=rep_weight_cols,
            method=method_str,
            fay_coef=final_fay,
            center=variance_center,
            degrees_of_freedom=df_val,
            by_col=by_col,
        )
        est_list = self._polars_result_to_param_est(
            result_df, y, PopParam.MEAN, alpha, deff=False, by_col=by_col, as_factor=False
        )
        variances = result_df["var"].to_numpy()
        est_cov = np.diag(variances)
        return self._build_estimate_result_light(
            est_list, est_cov, PopParam.MEAN, alpha, by_cols, as_factor=False, method=method
        )

    def _replicate_total_polars(
        self,
        y,
        *,
        method,
        by=None,
        fay_coef=0.0,
        variance_center="rep_mean",
        alpha=0.05,
        drop_nulls=False,
    ):
        design = self._sample._design
        _raw = self._sample._data
        data: pl.DataFrame = (
            cast(pl.DataFrame, _raw.collect())
            if isinstance(_raw, pl.LazyFrame)
            else cast(pl.DataFrame, _raw)
        )
        rw = design.rep_wgts
        if rw is None:
            raise ValueError("Replication weights required.")
        rep_weight_cols = self._get_rep_weight_cols()
        if not rep_weight_cols:
            raise ValueError("No replicate weight columns found.")
        n_reps = len(rep_weight_cols)
        df_val = int(rw.df) if rw.df and rw.df > 0 else max(1, n_reps - 1)
        final_fay = float(fay_coef) if fay_coef != 0.0 else float(rw.fay_coef)

        by_cols = _colspec_to_list(by)

        if len(by_cols) > 1:
            by_col = f"_by_{_INTERNAL_CONCAT_SUFFIX}"
            if by_col not in data.columns:
                expr = pl.concat_str(
                    [pl.col(c).cast(pl.String).fill_null("__Null__") for c in by_cols],
                    separator=_BY_SEP,
                )
                data = data.with_columns(expr.alias(by_col))
        elif len(by_cols) == 1:
            by_col = by_cols[0]
        else:
            by_col = None

        data = data.with_columns(pl.col(y).cast(pl.Float64))
        data = self._ensure_float64(data, rep_weight_cols)
        filter_cols = [y]
        weight_col = design.wgt
        if weight_col:
            filter_cols.append(weight_col)
        if by_col:
            filter_cols.append(by_col)
        if drop_nulls:
            data = data.filter(pl.all_horizontal(pl.col(c).is_not_null() for c in filter_cols))
        if by_col and data[by_col].dtype != pl.String:
            data = data.with_columns(pl.col(by_col).cast(pl.String))
        if not weight_col:
            weight_col = "__unit_wgt__"
            data = data.with_columns(pl.lit(1.0).alias(weight_col))
        else:
            data = self._ensure_float64(data, [weight_col])
        method_str = self._get_rep_method_str(method)
        result_df = rs.replicate_total(
            data,
            value_col=y,
            weight_col=weight_col,
            rep_weight_cols=rep_weight_cols,
            method=method_str,
            fay_coef=final_fay,
            center=variance_center,
            degrees_of_freedom=df_val,
            by_col=by_col,
        )
        est_list = self._polars_result_to_param_est(
            result_df, y, PopParam.TOTAL, alpha, deff=False, by_col=by_col, as_factor=False
        )
        variances = result_df["var"].to_numpy()
        est_cov = np.diag(variances)
        return self._build_estimate_result_light(
            est_list, est_cov, PopParam.TOTAL, alpha, by_cols, as_factor=False, method=method
        )

    def _replicate_ratio_polars(
        self,
        y,
        x,
        *,
        method,
        by=None,
        fay_coef=0.0,
        variance_center="rep_mean",
        alpha=0.05,
        drop_nulls=False,
    ):
        design = self._sample._design
        _raw = self._sample._data
        data: pl.DataFrame = (
            cast(pl.DataFrame, _raw.collect())
            if isinstance(_raw, pl.LazyFrame)
            else cast(pl.DataFrame, _raw)
        )
        rw = design.rep_wgts
        if rw is None:
            raise ValueError("Replication weights required.")
        rep_weight_cols = self._get_rep_weight_cols()
        if not rep_weight_cols:
            raise ValueError("No replicate weight columns found.")
        n_reps = len(rep_weight_cols)
        df_val = int(rw.df) if rw.df and rw.df > 0 else max(1, n_reps - 1)
        final_fay = float(fay_coef) if fay_coef != 0.0 else float(rw.fay_coef)

        by_cols = _colspec_to_list(by)

        if len(by_cols) > 1:
            by_col = f"_by_{_INTERNAL_CONCAT_SUFFIX}"
            if by_col not in data.columns:
                expr = pl.concat_str(
                    [pl.col(c).cast(pl.String).fill_null("__Null__") for c in by_cols],
                    separator=_BY_SEP,
                )
                data = data.with_columns(expr.alias(by_col))
        elif len(by_cols) == 1:
            by_col = by_cols[0]
        else:
            by_col = None

        data = data.with_columns(pl.col(y).cast(pl.Float64), pl.col(x).cast(pl.Float64))
        data = self._ensure_float64(data, rep_weight_cols)
        filter_cols = [y, x]
        weight_col = design.wgt
        if weight_col:
            filter_cols.append(weight_col)
        if by_col:
            filter_cols.append(by_col)
        if drop_nulls:
            data = data.filter(pl.all_horizontal(pl.col(c).is_not_null() for c in filter_cols))
        if by_col and data[by_col].dtype != pl.String:
            data = data.with_columns(pl.col(by_col).cast(pl.String))
        if not weight_col:
            weight_col = "__unit_wgt__"
            data = data.with_columns(pl.lit(1.0).alias(weight_col))
        else:
            data = self._ensure_float64(data, [weight_col])
        method_str = self._get_rep_method_str(method)
        result_df = rs.replicate_ratio(
            data,
            numerator_col=y,
            denominator_col=x,
            weight_col=weight_col,
            rep_weight_cols=rep_weight_cols,
            method=method_str,
            fay_coef=final_fay,
            center=variance_center,
            degrees_of_freedom=df_val,
            by_col=by_col,
        )
        est_list = self._polars_result_to_param_est(
            result_df,
            y,
            PopParam.RATIO,
            alpha,
            deff=False,
            by_col=by_col,
            as_factor=False,
            x_name=x,
        )
        variances = result_df["var"].to_numpy()
        est_cov = np.diag(variances)
        return self._build_estimate_result_light(
            est_list, est_cov, PopParam.RATIO, alpha, by_cols, as_factor=False, method=method
        )

    def _replicate_prop_polars(
        self,
        y,
        *,
        method,
        by=None,
        fay_coef=0.0,
        variance_center="rep_mean",
        alpha=0.05,
        ci_method="logit",
        drop_nulls=False,
    ):
        design = self._sample._design
        _raw = self._sample._data
        data: pl.DataFrame = (
            cast(pl.DataFrame, _raw.collect())
            if isinstance(_raw, pl.LazyFrame)
            else cast(pl.DataFrame, _raw)
        )
        data = data.with_columns(pl.col(y).cast(pl.Int64))
        rw = design.rep_wgts
        if rw is None:
            raise ValueError("Replication weights required.")
        rep_weight_cols = self._get_rep_weight_cols()
        if not rep_weight_cols:
            raise ValueError("No replicate weight columns found.")
        n_reps = len(rep_weight_cols)
        df_val = int(rw.df) if rw.df and rw.df > 0 else max(1, n_reps - 1)
        final_fay = float(fay_coef) if fay_coef != 0.0 else float(rw.fay_coef)

        by_cols = _colspec_to_list(by)

        if len(by_cols) > 1:
            by_col = f"_by_{_INTERNAL_CONCAT_SUFFIX}"
            if by_col not in data.columns:
                expr = pl.concat_str(
                    [pl.col(c).cast(pl.String).fill_null("__Null__") for c in by_cols],
                    separator=_BY_SEP,
                )
                data = data.with_columns(expr.alias(by_col))
        elif len(by_cols) == 1:
            by_col = by_cols[0]
        else:
            by_col = None

        data = self._ensure_float64(data, rep_weight_cols)
        filter_cols = [y]
        weight_col = design.wgt
        if weight_col:
            filter_cols.append(weight_col)
        if by_col:
            filter_cols.append(by_col)
        if drop_nulls:
            data = data.filter(pl.all_horizontal(pl.col(c).is_not_null() for c in filter_cols))
        if by_col and data[by_col].dtype != pl.String:
            data = data.with_columns(pl.col(by_col).cast(pl.String))
        if not weight_col:
            weight_col = "__unit_wgt__"
            data = data.with_columns(pl.lit(1.0).alias(weight_col))
        else:
            data = self._ensure_float64(data, [weight_col])
        method_str = self._get_rep_method_str(method)
        result_df = rs.replicate_prop(
            data,
            value_col=y,
            weight_col=weight_col,
            rep_weight_cols=rep_weight_cols,
            method=method_str,
            fay_coef=final_fay,
            center=variance_center,
            degrees_of_freedom=df_val,
            by_col=by_col,
        )
        est_list = self._polars_result_to_param_est(
            result_df,
            y,
            PopParam.PROP,
            alpha,
            deff=False,
            by_col=by_col,
            as_factor=True,
            ci_method=ci_method,
        )
        variances = result_df["var"].to_numpy()
        est_cov = np.diag(variances)
        return self._build_estimate_result_light(
            est_list, est_cov, PopParam.PROP, alpha, by_cols, as_factor=True, method=method
        )

    def _replicate_median_polars(
        self,
        y,
        *,
        method,
        by=None,
        fay_coef=0.0,
        q_method: QuantileMethod = QuantileMethod.HIGHER,
        variance_center="rep_mean",
        alpha=0.05,
        drop_nulls=False,
    ):
        """Replication-based median estimation using Rust backend."""
        design = self._sample._design
        _raw = self._sample._data
        data: pl.DataFrame = (
            cast(pl.DataFrame, _raw.collect())
            if isinstance(_raw, pl.LazyFrame)
            else cast(pl.DataFrame, _raw)
        )
        rw = design.rep_wgts
        if rw is None:
            raise ValueError("Replication weights required for replication-based estimation.")
        rep_weight_cols = self._get_rep_weight_cols()
        if not rep_weight_cols:
            raise ValueError("No replicate weight columns found.")
        n_reps = len(rep_weight_cols)
        df_val = int(rw.df) if rw.df and rw.df > 0 else max(1, n_reps - 1)
        final_fay = float(fay_coef) if fay_coef != 0.0 else float(rw.fay_coef)

        by_cols = _colspec_to_list(by)

        if len(by_cols) > 1:
            by_col = f"_by_{_INTERNAL_CONCAT_SUFFIX}"
            if by_col not in data.columns:
                expr = pl.concat_str(
                    [pl.col(c).cast(pl.String).fill_null("__Null__") for c in by_cols],
                    separator=_BY_SEP,
                )
                data = data.with_columns(expr.alias(by_col))
        elif len(by_cols) == 1:
            by_col = by_cols[0]
        else:
            by_col = None

        data = data.with_columns(pl.col(y).cast(pl.Float64))
        data = self._ensure_float64(data, rep_weight_cols)
        filter_cols = [y]
        weight_col = design.wgt
        if weight_col:
            filter_cols.append(weight_col)
        if by_col:
            filter_cols.append(by_col)
        if drop_nulls:
            data = data.filter(pl.all_horizontal(pl.col(c).is_not_null() for c in filter_cols))
        if by_col and data[by_col].dtype != pl.String:
            data = data.with_columns(pl.col(by_col).cast(pl.String))
        if not weight_col:
            weight_col = "__unit_wgt__"
            data = data.with_columns(pl.lit(1.0).alias(weight_col))
        else:
            data = self._ensure_float64(data, [weight_col])

        method_str = self._get_rep_method_str(method)
        q_method_str = q_method.value if hasattr(q_method, "value") else str(q_method).lower()

        result_df = rs.replicate_median(
            data,
            value_col=y,
            weight_col=weight_col,
            rep_weight_cols=rep_weight_cols,
            method=method_str,
            fay_coef=final_fay,
            center=variance_center,
            degrees_of_freedom=df_val,
            by_col=by_col,
            quantile_method=q_method_str,
        )

        est_list = self._replicate_median_result_to_param_est(result_df, y, alpha, by_col)
        variances = result_df["var"].to_numpy()
        est_cov = np.diag(variances)

        return self._build_estimate_result_light(
            est_list, est_cov, PopParam.MEDIAN, alpha, by_cols, as_factor=False, method=method
        )
