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

import copy
import logging
import re

from collections.abc import Iterable, Mapping
from typing import Self, cast

import msgspec
import numpy as np

from svy.core.enumerations import OnePropSizeMethod, PopParam, PropVarMode, TwoPropsSizeMethod
from svy.core.types import DomainScalarMap, Number
from svy.engine.size_and_power.size import (
    _apply_deff,
    _apply_deff_pair,
    _apply_fpc_srswor,
    _apply_fpc_srswor_pair,
    _apply_nonresponse,
    _apply_nonresponse_pair,
    _fleiss_sample_size_prop,
    _wald_sample_size_mean,
    _wald_sample_size_prop,
    _wald_sample_size_two_props,
)
from svy.utils.helpers import _get_keys_from_maps


log = logging.getLogger(__name__)

_CHUNK_RE = re.compile(r"(\d+)")


# =============================================================================
# Optional central UI integration (svy.ui.printing)
# =============================================================================


def _ui_module():
    """Return svy.ui.printing if available, else None (soft dependency)."""
    try:
        from svy.ui import printing as _p

        return _p
    except Exception:
        return None


def _resolve_width_for(obj, default: int = 88) -> int:
    """
    Ask svy.ui.printing for a preferred width when available; otherwise use a
    robust fallback chain (instance → class → module var → __main__ → builtins
    → env → terminal → default).
    """
    P = _ui_module()
    if P:
        for call in (
            lambda: P.resolve_width(owner=obj, default=default),
            lambda: P.resolve_width(default=default),
            lambda: P.resolve_width(),
        ):
            try:
                w = int(call())
                if w > 20:
                    return w
            except Exception:
                pass

    import builtins
    import os
    import shutil
    import sys

    def _as_int(x):
        try:
            return int(x)
        except Exception:
            return None

    w = _as_int(getattr(obj, "_print_width", None))
    if isinstance(w, int) and w > 20:
        return w

    w = _as_int(getattr(type(obj), "PRINT_WIDTH", None))
    if isinstance(w, int) and w > 20:
        return w

    try:
        mod = sys.modules.get(obj.__class__.__module__)
        if mod is not None and hasattr(mod, "SVY_PRINT_WIDTH"):
            mw = _as_int(getattr(mod, "SVY_PRINT_WIDTH"))
            if isinstance(mw, int) and mw > 20:
                return mw
    except Exception:
        pass

    main_mod = sys.modules.get("__main__")
    if main_mod is not None and hasattr(main_mod, "SVY_PRINT_WIDTH"):
        mw = _as_int(getattr(main_mod, "SVY_PRINT_WIDTH"))
        if isinstance(mw, int) and mw > 20:
            return mw

    bw = _as_int(getattr(builtins, "SVY_PRINT_WIDTH", None))
    if isinstance(bw, int) and bw > 20:
        return bw

    ew = _as_int(os.environ.get("SVY_PRINT_WIDTH"))
    if isinstance(ew, int) and ew > 20:
        return ew

    try:
        ts = shutil.get_terminal_size(fallback=(default, 24))
        if isinstance(ts.columns, int) and ts.columns > 20:
            return ts.columns
    except Exception:
        pass

    return default


def _panel_settings_for(
    obj,
    *,
    title_hint: str = "Sample Size",
    border_fallback: str = "cyan",
    feature_key: str = "sampsize.results",
) -> tuple[bool, str, str]:
    """
    Resolve (enabled, title, border_color) for an outer panel. Uses central UI
    if available; falls back to class vars, then env override.
    """
    show_panel = bool(getattr(type(obj), "_PRINT_PANEL", True))
    title = getattr(type(obj), "_PANEL_TITLE", title_hint)
    border = getattr(type(obj), "_PANEL_BORDER", border_fallback)

    P = _ui_module()
    if P:
        try:
            enabled = P.panel_enabled(feature_key)
            if isinstance(enabled, bool):
                show_panel = enabled
        except Exception:
            pass
        try:
            st = P.styles(feature_key, title=title, border=border)
            if isinstance(st, dict):
                title = st.get("title", title)
                border = st.get("border", border)
        except Exception:
            pass

    import os as _os

    env_panel = _os.environ.get("SVY_PRINT_PANEL")
    if env_panel is not None:
        show_panel = env_panel.lower() not in {"0", "false", "off", "no"}

    return show_panel, title, border


# =============================================================================
# Core data structures
# =============================================================================


class SampleSizeConfigs(msgspec.Struct):
    stratified: bool = False
    pop_size: Number | DomainScalarMap | None = None
    deff: Number | DomainScalarMap = 1.0
    resp_rate: Number | DomainScalarMap = 1.0
    alloc_unit: Number | DomainScalarMap = 1  # round up to multiple (e.g., cluster size)


class TargetProp(msgspec.Struct, frozen=True, tag="prop"):
    p: Number
    moe: Number
    alpha: Number = 0.05


class TargetMean(msgspec.Struct, frozen=True, tag="mean"):
    sigma: Number
    moe: Number
    alpha: Number = 0.05


class TargetTwoProps(msgspec.Struct, frozen=True, tag="two_prop"):
    p1: Number
    p2: Number
    alloc_ratio: Number = 1.0  # allocation ratio n2/n1
    alpha: Number = 0.05
    power: Number = 0.80


class TargetTwoMeans(msgspec.Struct, frozen=True, tag="two_prop"):
    mu1: Number
    mu2: Number
    alloc_ratio: Number = 1.0  # allocation ratio n2/n1
    alpha: Number = 0.05
    power: Number = 0.80


Target = TargetProp | TargetMean | TargetTwoProps | TargetTwoMeans


class Size(msgspec.Struct, frozen=True, tag="size"):
    stratum: str | None = None
    n0: Number | tuple[Number, Number] = 0  # base (no FPC/DEFF/response)
    n1_fpc: Number | tuple[Number, Number] | None = None  # after FPC (if pop_size provided)
    n2_deff: Number | tuple[Number, Number] | None = None  # after DEFF (if response provided)
    n: Number | tuple[Number, Number] = (
        0  # final rounded after nonresponse adjustment; for two-arm: (n1, n2)
    )


class SampleSize:
    """
    Compute required sample sizes for survey objectives under simple or stratified designs.

    Parameters
    ----------
    pop_size : Number | DomainScalarMap | None, default None
        Target population size. When provided as a mapping, values are interpreted
        per stratum (keys are stratum labels). If omitted, no finite population
        correction is applied.
    deff : Number | DomainScalarMap, default 1.0
        Design effect. May be a scalar or a per-stratum mapping.
    resp_rate : Number | DomainScalarMap, default 1.0
        Anticipated unit response rate. May be a scalar or a per-stratum mapping.
    alloc_unit : Number | DomainScalarMap, default 1
        Allocation unit (e.g., cluster/PSU size or minimum take). May be a scalar
        or a per-stratum mapping.
    stratified : bool, default False
        If ``True``, interpret inputs as stratified. See notes below for rules.
    n_strata : int | None, optional
        Number of strata when ``stratified=True`` and all other parameters are
        scalar; values will be replicated across strata.

    Notes
    -----
    - **Unstratified**: all of ``pop_size``, ``deff``, ``resp_rate``, and
      ``alloc_unit`` must be scalars.
    - **Stratified (per-stratum mappings provided)**: pass dicts for any of the
      parameters to supply per-stratum values.
    - **Stratified (all scalars)**: set ``n_strata`` to a positive integer; scalar
      values will be replicated across strata.
    """

    # ---------- printing/branding toggles (class-level fallbacks) ----------
    _PRINT_PANEL: bool = True
    _PANEL_TITLE: str = "Sample Size"
    _PANEL_BORDER: str = "cyan"
    PRINT_WIDTH: int | None = None  # optional class-level width override

    __slots__ = ("_configs", "_param", "_target", "_size")

    def __init__(
        self,
        pop_size: Number | DomainScalarMap | None = None,
        deff: Number | DomainScalarMap = 1.0,
        resp_rate: Number | DomainScalarMap = 1.0,
        alloc_unit: Number | DomainScalarMap = 1,
        stratified: bool = False,
        n_strata: Number | None = None,
    ):
        at_least_one_dict = (
            isinstance(pop_size, dict)
            or isinstance(deff, dict)
            or isinstance(resp_rate, dict)
            or isinstance(alloc_unit, dict)
        )

        if not stratified and at_least_one_dict:
            raise ValueError(
                "pop_size, deff, resp_rate, and alloc_unit must be scalar when stratified=False"
            )
        elif not stratified:
            self._configs = SampleSizeConfigs(
                stratified=False,
                pop_size=pop_size,
                deff=deff,
                resp_rate=resp_rate,
                alloc_unit=alloc_unit,
            )
        else:
            if (n_strata is None or n_strata <= 0) and not at_least_one_dict:
                raise ValueError(
                    "n_strata must be a positive integer when stratified=True and all other parameters are scalar"
                )
            elif at_least_one_dict:
                strata = _get_keys_from_maps(
                    **{
                        k: v
                        for k, v in dict(
                            pop_size=pop_size,
                            deff=deff,
                            resp_rate=resp_rate,
                            alloc_unit=alloc_unit,
                        ).items()
                        if isinstance(v, Mapping)
                    }
                )
                pop_size_map = (
                    {k: pop_size for k in strata} if isinstance(pop_size, Number) else pop_size
                )
                deff_map = {k: deff for k in strata} if isinstance(deff, Number) else deff
                resp_rate_map = (
                    {k: resp_rate for k in strata} if isinstance(resp_rate, Number) else resp_rate
                )
                alloc_unit_map = (
                    {k: alloc_unit for k in strata}
                    if isinstance(alloc_unit, Number)
                    else alloc_unit
                )

                self._configs = SampleSizeConfigs(
                    stratified=True,
                    pop_size=pop_size_map,
                    deff=deff_map,
                    resp_rate=resp_rate_map,
                    alloc_unit=alloc_unit_map,
                )
            else:
                assert n_strata is not None and n_strata > 0
                pop_size_dict = {}
                deff_dict = {}
                resp_rate_dict = {}
                alloc_unit_dict = {}
                for i in range(int(n_strata)):
                    pop_size_dict[i] = pop_size
                    deff_dict[i] = deff
                    resp_rate_dict[i] = resp_rate
                    alloc_unit_dict[i] = alloc_unit
                self._configs = SampleSizeConfigs(
                    stratified=True,
                    pop_size=pop_size_dict,
                    deff=deff_dict,
                    resp_rate=resp_rate_dict,
                    alloc_unit=alloc_unit_dict,
                )

        self._param = None
        self._target = None
        self._size = None

    # ---------- internal helpers ----------
    @staticmethod
    def _is_size_obj(obj) -> bool:
        return all(hasattr(obj, a) for a in ("n0", "n1_fpc", "n2_deff", "n"))

    def _iter_sizes(self):
        """Return a flat list of Size-like objects from self.size/_size."""
        sizes = getattr(self, "size", None)
        if sizes is None:
            sizes = getattr(self, "_size", None)
        if sizes is None:
            return []

        if self._is_size_obj(sizes):  # single Size
            return [sizes]

        if isinstance(sizes, Mapping):  # dict of Size or dict -> [Size]
            out = []
            for v in sizes.values():
                if self._is_size_obj(v):
                    out.append(v)
                elif isinstance(v, Iterable) and not isinstance(v, (str, bytes)):
                    out.extend([x for x in v if self._is_size_obj(x)])
            return out

        if isinstance(sizes, Iterable) and not isinstance(sizes, (str, bytes)):
            return [v for v in sizes if self._is_size_obj(v)]

        return []

    @staticmethod
    def _fmt_num(x) -> str:
        try:
            xf = float(x)
            xi = round(xf)
            return f"{xi:.0f}" if abs(xf - xi) < 1e-9 else f"{xf:.1f}"
        except Exception:
            return str(x)

    @staticmethod
    def _nat_key(label: object):
        """Natural sort: 'region2' < 'region10'; None/'overall' first."""
        if label is None:
            return (0, [""])
        s = str(label).strip().lower()
        is_overall = 0 if s in {"", "overall", "all", "national"} else 1
        parts = _CHUNK_RE.split(s)
        key = [int(p) if p.isdigit() else p for p in parts]
        return (is_overall, key)

    # ---------------- public: Polars export -------------------
    def to_polars(
        self,
        order_by: str | None = None,
        ascending: bool = True,
        natural: bool = True,
        overall_first: bool = True,  # kept for API compatibility
        group_labels: list[str] | None = None,
        include_stratum: bool | None = None,
    ):
        """
        Normalize sizes to a Polars DataFrame.

        - If `include_stratum` is None, it defaults to `self._configs.stratified`.
        - When `include_stratum` is False, the resulting DF will NOT have a 'stratum' column.
        - Handles scalar, stratified, and comparison (tuple) sizes.
        """
        try:
            import polars as pl
        except Exception as e:
            raise ImportError("to_polars() requires the 'polars' package") from e

        if include_stratum is None:
            include_stratum = bool(
                getattr(self, "_configs", None) and getattr(self._configs, "stratified", False)
            )

        objs = self._iter_sizes()
        if not objs:
            cols = (["stratum"] if include_stratum else []) + ["n0", "n1_fpc", "n2_deff", "n"]
            return pl.DataFrame({c: [] for c in cols})

        # detect tuple (comparison) case and validate lengths
        from collections.abc import Iterable as _Iterable

        def _len_if_iter(x):
            return len(x) if isinstance(x, _Iterable) and not isinstance(x, (str, bytes)) else None

        fields = ("n0", "n1_fpc", "n2_deff", "n")
        group_len = None
        for s in objs:
            lens = {f: _len_if_iter(getattr(s, f)) for f in fields}
            ls = [L for L in lens.values() if L is not None]
            if ls:
                L0 = ls[0]
                if any(L != L0 for L in ls):
                    raise ValueError(f"Inconsistent tuple lengths within a Size: {lens}")
                if group_len is None:
                    group_len = L0
                elif group_len != L0:
                    raise ValueError(
                        f"Inconsistent tuple lengths across strata: expected {group_len}, got {L0}"
                    )

        rows = []
        if group_len is None:
            # scalar case
            for s in objs:
                label = getattr(s, "stratum", None) or getattr(s, "domain", None) or "overall"
                row: dict[str, object] = {
                    "n0": float(s.n0),
                    "n1_fpc": float(s.n1_fpc),
                    "n2_deff": float(s.n2_deff),
                    "n": float(s.n),
                }
                if include_stratum:
                    row["stratum"] = "" if label is None else str(label)
                rows.append(row)
            df = pl.DataFrame(rows)
            cols = (["stratum"] if include_stratum else []) + ["n0", "n1_fpc", "n2_deff", "n"]
            df = df.select(cols)
        else:
            # tuple (comparison) case
            if not group_labels:
                group_labels = [f"g{i + 1}" for i in range(group_len)]
            if len(group_labels) != group_len:
                raise ValueError(f"group_labels length {len(group_labels)} != {group_len}")

            for s in objs:
                label = getattr(s, "stratum", None) or getattr(s, "domain", None) or "overall"
                n0, n1, n2, n = (getattr(s, f) for f in fields)
                for i in range(group_len):
                    row: dict[str, object] = {
                        "group": group_labels[i],
                        "n0": float(n0[i]),
                        "n1_fpc": float(n1[i]),
                        "n2_deff": float(n2[i]),
                        "n": float(n[i]),
                    }
                    if include_stratum:
                        row["stratum"] = "" if label is None else str(label)
                    rows.append(row)
            df = pl.DataFrame(rows)
            cols = (["stratum"] if include_stratum else []) + [
                "group",
                "n0",
                "n1_fpc",
                "n2_deff",
                "n",
            ]
            df = df.select(cols)

        # ----- sorting defaults / behavior
        if order_by is None:
            if include_stratum:
                order_by = "stratum"
            else:
                order_by = "group" if ("group" in df.columns) else "n"

        if order_by == "stratum" and include_stratum and "stratum" in df.columns:
            if natural:
                lower = df["stratum"].str.to_lowercase()
                is_overall = lower.is_in(["", "overall", "all", "national"]).cast(pl.Int8)
                # overall first (1 - is_overall) → 0 for overall; reverse if overall_first=False
                overall_key = (1 - is_overall) if overall_first else is_overall
                df = (
                    df.with_columns(
                        [
                            overall_key.alias("_overall"),
                            lower.str.replace_all(r"\d+", "").alias("_prefix"),
                            lower.str.extract(r"(\d+)").cast(pl.Int64).fill_null(0).alias("_num"),
                        ]
                    )
                    .sort(["_overall", "_prefix", "_num"], descending=[not ascending] * 3)
                    .drop(["_overall", "_prefix", "_num"])
                )
            else:
                df = df.sort("stratum", descending=not ascending)

        elif order_by == "group" and "group" in df.columns:
            if natural:
                gl = df["group"].str.to_lowercase()
                df = (
                    df.with_columns(
                        [
                            gl.str.replace_all(r"\d+", "").alias("_gpre"),
                            gl.str.extract(r"(\d+)").cast(pl.Int64).fill_null(0).alias("_gnum"),
                        ]
                    )
                    .sort(["_gpre", "_gnum"], descending=[not ascending, not ascending])
                    .drop(["_gpre", "_gnum"])
                )
            else:
                df = df.sort("group", descending=not ascending)

        elif order_by in {"n", "n0", "n1_fpc", "n2_deff"}:
            df = df.sort(order_by, descending=not ascending)
        else:
            pass

        # Stable secondary sort when both columns exist
        if include_stratum and "group" in df.columns and order_by in {"stratum", "group"}:
            gl = df["group"].str.to_lowercase() if "group" in df.columns else None
            if gl is not None:
                df = (
                    df.with_columns(
                        [
                            df["stratum"]
                            .str.to_lowercase()
                            .str.replace_all(r"\d+", "")
                            .alias("_spre"),
                            df["stratum"]
                            .str.to_lowercase()
                            .str.extract(r"(\d+)")
                            .cast(pl.Int64)
                            .fill_null(0)
                            .alias("_snum"),
                            gl.str.replace_all(r"\d+", "").alias("_gpre2"),
                            gl.str.extract(r"(\d+)").cast(pl.Int64).fill_null(0).alias("_gnum2"),
                        ]
                    )
                    .sort(
                        ["_spre", "_snum", "_gpre2", "_gnum2"],
                        descending=[False, False, False, False],
                    )
                    .drop(["_spre", "_snum", "_gpre2", "_gnum2"])
                )

        return df

    def _format_table_ascii(self) -> str:
        """ASCII fallback that hides 'stratum' when not stratified."""
        objs = self._iter_sizes()
        if not objs:
            return "<no sizes>"

        stratified = bool(
            getattr(self, "_configs", None) and getattr(self._configs, "stratified", False)
        )

        from collections.abc import Iterable as _Iterable

        any_tuple = any(
            isinstance(getattr(s, "n"), _Iterable)
            and not isinstance(getattr(s, "n"), (str, bytes))
            for s in objs
        )

        # Build rows according to shape & stratified flag
        if any_tuple:
            headers = (
                ("group", "n0", "n1_fpc", "n2_deff", "n")
                if not stratified
                else ("stratum", "group", "n0", "n1_fpc", "n2_deff", "n")
            )
            rows = []
            for s in objs:
                label = getattr(s, "stratum", None) or getattr(s, "domain", None) or "overall"
                n0, n1, n2, n = s.n0, s.n1_fpc, s.n2_deff, s.n
                m = len(n)
                for i in range(m):
                    if stratified:
                        rows.append((str(label), f"g{i + 1}", n0[i], n1[i], n2[i], n[i]))
                    else:
                        rows.append((f"g{i + 1}", n0[i], n1[i], n2[i], n[i]))
            # natural sort
            if stratified:
                rows.sort(key=lambda r: (self._nat_key(r[0]), self._nat_key(r[1])))
            else:
                rows.sort(key=lambda r: self._nat_key(r[0]))
        else:
            headers = (
                ("n0", "n1_fpc", "n2_deff", "n")
                if not stratified
                else ("stratum", "n0", "n1_fpc", "n2_deff", "n")
            )
            rows = []
            for s in objs:
                label = getattr(s, "stratum", None) or getattr(s, "domain", None) or "overall"
                if stratified:
                    rows.append((str(label), s.n0, s.n1_fpc, s.n2_deff, s.n))
                else:
                    rows.append((s.n0, s.n1_fpc, s.n2_deff, s.n))

            if stratified:
                rows.sort(key=lambda r: self._nat_key(r[0]))

        # format numbers
        def fmt_row(r):
            # left text cols stay raw; numeric cols formatted
            if len(headers) == 6:  # stratified + groups
                return (r[0], r[1], *(self._fmt_num(v) for v in r[2:]))
            elif len(headers) == 5:  # stratified scalar
                return (r[0], *(self._fmt_num(v) for v in r[1:]))
            elif len(headers) == 4:  # non-stratified scalar
                return tuple(self._fmt_num(v) for v in r)
            else:  # non-stratified groups (5 cols)
                return (r[0], *(self._fmt_num(v) for v in r[1:]))

        formatted = [tuple(map(str, fmt_row(r))) for r in rows]

        # compute widths
        cols = list(zip(*([headers] + formatted)))
        widths = [max(len(c) for c in col) for col in cols]

        def line(cells):
            out = []
            for i, (c, w) in enumerate(zip(cells, widths)):
                left_aligned = headers[i] in {"stratum", "group"}
                out.append(c.ljust(w) if left_aligned else c.rjust(w))
            return "  ".join(out)

        out = [line(headers), line(["-" * w for w in widths])]
        out += [line(r) for r in formatted]
        return "\n".join(out)

    # ---------- indent/surround helper ----------
    @staticmethod
    def _pad_and_surround(text: str, *, indent: int = 2, surround: bool = False) -> str:
        text = str(text).rstrip("\n")
        if indent > 0:
            pad = " " * indent
            text = "\n".join(pad + line if line else pad for line in text.splitlines())
        if surround:
            return f"\n{text}\n"
        return text

    # ----------------- width-aware, panel-styled string -------------------
    def __plain_str__(self) -> str:
        """Plain-text fallback when rich is not installed. Never calls str(self)."""
        ascii_tbl = self._format_table_ascii()
        return self._pad_and_surround(ascii_tbl, indent=2, surround=False)

    def __str__(self) -> str:
        from svy.ui.printing import render_rich_to_str

        return render_rich_to_str(self, width=_resolve_width_for(self))

    def __repr__(self) -> str:
        n = len(self._iter_sizes())
        stratified = getattr(self._configs, "stratified", False)
        param = getattr(self._param, "name", None) or "uncomputed"
        return f"SampleSize(stratified={stratified}, strata={n}, param={param})"

    # ---------------- Rich integration (optional) --------------
    def to_rich_table(self, title: str | None = None):
        """Return a Rich Table (requires `rich`)."""
        try:
            from rich import box
            from rich.table import Table
        except Exception as e:
            raise ImportError("`rich` is required for to_rich_table(). `pip install rich`") from e

        # Try to use Polars for normalized rows
        cols, rows = None, None
        try:
            df = self.to_polars()
            cols, rows = df.columns, df.rows()
            start_idx = 2 if ("group" in cols) else 1
            rows = [
                tuple((v if i < start_idx else self._fmt_num(v)) for i, v in enumerate(row))
                for row in rows
            ]
        except Exception:
            # fallback to ASCII normalization
            from rich.text import Text

            t = Table(title=title, box=box.SIMPLE)
            t.add_column("output")
            t.add_row(Text(self._format_table_ascii()))
            return t

        table = Table(
            title=title, box=box.SIMPLE, show_header=True, header_style="bold", pad_edge=False
        )
        for i, c in enumerate(cols):
            justify = "left" if c in ("stratum", "group") else "right"
            table.add_column(c, justify=justify, overflow="fold")
        for r in rows:
            table.add_row(*[str(x) for x in r])
        return table

    def __rich_console__(self, console, options):
        try:
            from rich.console import Group
            from rich.padding import Padding
            from rich.panel import Panel

            tbl = self.to_rich_table()
            content = Group(Padding(tbl, (0, 0, 0, 2)))

            show_panel, panel_title, border = _panel_settings_for(
                self,
                title_hint=getattr(type(self), "_PANEL_TITLE", "Sample Size"),
                border_fallback=getattr(type(self), "_PANEL_BORDER", "cyan"),
                feature_key="sampsize.results",
            )

            if show_panel:
                yield Panel(content, title=panel_title, border_style=border)
            else:
                yield content
        except Exception:
            from rich.text import Text

            yield Text(self.__plain_str__())

    ## PROPERTY METHODS

    @property
    def configs(self) -> SampleSizeConfigs:
        return self._configs

    @property
    def target(self) -> Target | None:
        return self._target

    @property
    def size(self) -> Size | list[Size] | None:
        return self._size

    @property
    def param(self) -> PopParam | None:
        """Return a defensive copy to avoid external mutation."""
        return copy.copy(self._param)

    ## SAMPLE SIZES FOR PROPORTIONS

    def estimate_prop(
        self,
        p: Number | DomainScalarMap,
        moe: Number | DomainScalarMap,
        *,
        method: OnePropSizeMethod | DomainScalarMap = OnePropSizeMethod.WALD,
        alpha: Number | DomainScalarMap = 0.05,
    ) -> Self:
        self._param = PopParam.PROP

        at_least_one_dict = (
            isinstance(p, dict)
            or isinstance(moe, dict)
            or isinstance(method, dict)
            or isinstance(alpha, dict)
        )

        if self._configs.stratified and not at_least_one_dict:
            raise ValueError(
                "At least one of p, moe, method, or alpha must be a dictionary when stratified sampling is used."
            )
        elif self._configs.stratified:
            strata = _get_keys_from_maps(
                **{
                    k: v
                    for k, v in dict(p=p, moe=moe, method=method, alpha=alpha).items()
                    if isinstance(v, Mapping)
                }
            )
            m = len(strata)
            params = {"p": p, "moe": moe, "method": method, "alpha": alpha}
            for k, v in list(params.items()):
                if isinstance(v, (int, float, np.integer, np.floating)) and not isinstance(
                    v, bool
                ):
                    params[k] = dict(zip(strata, [v] * m))
            p, moe, method, alpha = params["p"], params["moe"], params["method"], params["alpha"]
        else:
            all_scalars = all(isinstance(x, (int, float)) for x in [p, moe, alpha])
            assert all_scalars, "All inputs must be scalars when not stratified."

        match method:
            case OnePropSizeMethod.WALD:
                n0 = _wald_sample_size_prop(
                    target=p,
                    half_ci=moe,
                    alpha=alpha,
                )
            case OnePropSizeMethod.FLEISS:
                n0 = _fleiss_sample_size_prop(
                    target=p,
                    half_ci=moe,
                    alpha=alpha,
                )
            case _:
                raise ValueError(f"Invalid method: {method}")

        if self._configs.pop_size:
            n1_fpc = _apply_fpc_srswor(n0=n0, pop_size=self._configs.pop_size)
        else:
            n1_fpc = n0
        n2_deff = _apply_deff(n=n1_fpc, deff=self._configs.deff)
        n_final = _apply_nonresponse(n=n2_deff, resp_rate=self._configs.resp_rate)

        if not self._configs.stratified:
            self._size = Size(
                stratum=None,
                n0=cast(Number, n0),
                n1_fpc=cast(Number, n1_fpc),
                n2_deff=cast(Number, n2_deff),
                n=cast(Number, n_final),
            )
        else:
            _n0 = cast(DomainScalarMap, n0)
            _n1_fpc = cast(DomainScalarMap, n1_fpc)
            _n2_deff = cast(DomainScalarMap, n2_deff)
            _n_final = cast(DomainScalarMap, n_final)
            self._size = [
                Size(
                    stratum=str(s),
                    n0=_n0[s],
                    n1_fpc=_n1_fpc[s],
                    n2_deff=_n2_deff[s],
                    n=_n_final[s],
                )
                for s in _n0
            ]

        return self

    def compare_props(
        self,
        p1: Number | DomainScalarMap,
        p2: Number | DomainScalarMap,
        *,
        two_sides: bool = True,
        delta: Number | DomainScalarMap = 0.0,
        alloc_ratio: Number = 1.0,  # allocation ratio n2/n1
        method: TwoPropsSizeMethod = TwoPropsSizeMethod.WALD,
        alpha: Number | DomainScalarMap = 0.05,
        power: Number | DomainScalarMap = 0.80,
        var_mode: PropVarMode = PropVarMode.ALT_PROPS,
    ) -> Self:
        self._param = PopParam.PROP

        at_least_one_dict = (
            isinstance(p1, dict)
            or isinstance(p2, dict)
            or isinstance(delta, dict)
            or isinstance(alloc_ratio, dict)
            or isinstance(alpha, dict)
            or isinstance(power, dict)
        )

        if self._configs.stratified and not at_least_one_dict:
            raise ValueError(
                "At least one of p, moe, method, or alpha must be a dictionary when stratified sampling is used."
            )
        elif self._configs.stratified:
            strata = _get_keys_from_maps(
                **{
                    k: v
                    for k, v in dict(
                        p1=p1,
                        p2=p2,
                        delta=delta,
                        alloc_ratio=alloc_ratio,
                        alpha=alpha,
                        power=power,
                    ).items()
                    if isinstance(v, Mapping)
                }
            )
            m = len(strata)
            params = {
                "p1": p1,
                "p2": p2,
                "delta": delta,
                "alloc_ratio": alloc_ratio,
                "alpha": alpha,
                "power": power,
            }
            for k, v in list(params.items()):
                if isinstance(v, (int, float, np.integer, np.floating)) and not isinstance(
                    v, bool
                ):
                    params[k] = dict(zip(strata, [v] * m))

            p1, p2, delta, alloc_ratio, alpha, power = (
                params["p1"],
                params["p2"],
                params["delta"],
                params["alloc_ratio"],
                params["alpha"],
                params["power"],
            )
        else:
            all_scalars = all(
                isinstance(x, (int, float)) for x in [p1, p2, delta, alloc_ratio, alpha, power]
            )
            assert all_scalars, "All inputs must be scalars when not stratified."

        if self._configs.stratified:
            _p1_map = cast(DomainScalarMap, p1)
            _p2_map = cast(DomainScalarMap, p2)
            epsilon: DomainScalarMap = {}
            for s in _p1_map:
                epsilon[s] = _p2_map[s] - _p1_map[s]
        else:
            epsilon = cast(Number, p2) - cast(Number, p1)

        match method:
            case TwoPropsSizeMethod.WALD:
                n0 = _wald_sample_size_two_props(
                    two_sides=two_sides,
                    epsilon=epsilon,
                    prop_1=p1,
                    prop_2=p2,
                    delta=delta,
                    kappa=alloc_ratio,
                    alpha=alpha,
                    power=power,
                    var_mode=var_mode,
                )
            case TwoPropsSizeMethod.MIETTINEN_NURMINEN:
                raise NotImplementedError("MIETTINEN_NURMINEN method is not implemented yet.")
            case TwoPropsSizeMethod.NEWCOMBE:
                raise NotImplementedError("NEWCOMBE method is not implemented yet.")
            case TwoPropsSizeMethod.FARRINGTON_MANNING:
                raise NotImplementedError("FARRINGTON_MANNING method is not implemented yet.")

        if self._configs.pop_size:
            n1_fpc = _apply_fpc_srswor_pair(n0=n0, pop_size=self._configs.pop_size)
        else:
            n1_fpc = n0

        n2_deff = _apply_deff_pair(n=n1_fpc, deff=self._configs.deff)
        n_final = _apply_nonresponse_pair(n=n2_deff, resp_rate=self._configs.resp_rate)

        if not self._configs.stratified:
            self._size = Size(
                stratum=None,
                n0=cast(Number, n0),
                n1_fpc=cast(Number, n1_fpc),
                n2_deff=cast(Number, n2_deff),
                n=cast(Number, n_final),
            )
        else:
            _n0 = cast(DomainScalarMap, n0)
            _n1_fpc = cast(DomainScalarMap, n1_fpc)
            _n2_deff = cast(DomainScalarMap, n2_deff)
            _n_final = cast(DomainScalarMap, n_final)
            self._size = [
                Size(
                    stratum=str(s),
                    n0=_n0[s],
                    n1_fpc=_n1_fpc[s],
                    n2_deff=_n2_deff[s],
                    n=_n_final[s],
                )
                for s in _n0
            ]

        return self

    ## SAMPLE SIZES FOR MEANS

    def estimate_mean(
        self,
        sigma: Number | DomainScalarMap,
        moe: Number | DomainScalarMap,
        *,
        method: OnePropSizeMethod | DomainScalarMap = OnePropSizeMethod.WALD,
        alpha: Number = 0.05,
    ) -> Self:
        self._param = PopParam.MEAN

        at_least_one_dict = (
            isinstance(sigma, dict)
            or isinstance(moe, dict)
            or isinstance(method, dict)
            or isinstance(alpha, dict)
        )

        if self._configs.stratified and not at_least_one_dict:
            raise ValueError(
                "At least one of p, moe, method, or alpha must be a dictionary when stratified sampling is used."
            )
        elif self._configs.stratified:
            strata = _get_keys_from_maps(
                **{
                    k: v
                    for k, v in dict(sigma=sigma, moe=moe, method=method, alpha=alpha).items()
                    if isinstance(v, Mapping)
                }
            )
            m = len(strata)
            params = {"sigma": sigma, "moe": moe, "method": method, "alpha": alpha}
            for k, v in list(params.items()):
                if isinstance(v, (int, float, np.integer, np.floating)) and not isinstance(
                    v, bool
                ):
                    params[k] = dict(zip(strata, [v] * m))

            sigma, moe, method, alpha = (
                params["sigma"],
                params["moe"],
                params["method"],
                params["alpha"],
            )
        else:
            all_scalars = all(isinstance(x, (int, float)) for x in [sigma, moe, alpha])
            assert all_scalars, "All inputs must be scalars when not stratified."

        match method:
            case OnePropSizeMethod.WALD:
                n0 = _wald_sample_size_mean(
                    half_ci=moe,
                    sigma=sigma,
                    alpha=alpha,
                )
            case OnePropSizeMethod.FLEISS:
                raise NotImplementedError("Fleiss method is not implemented for the mean.")
            case _:
                raise ValueError(f"Invalid method: {method}")

        if self._configs.pop_size:
            n1_fpc = _apply_fpc_srswor(n0=n0, pop_size=self._configs.pop_size)
        else:
            n1_fpc = n0

        n2_deff = _apply_deff(n=n1_fpc, deff=self._configs.deff)
        n_final = _apply_nonresponse(n=n2_deff, resp_rate=self._configs.resp_rate)

        if not self._configs.stratified:
            self._size = Size(
                stratum=None,
                n0=cast(Number, n0),
                n1_fpc=cast(Number, n1_fpc),
                n2_deff=cast(Number, n2_deff),
                n=cast(Number, n_final),
            )
        else:
            _n0 = cast(DomainScalarMap, n0)
            _n1_fpc = cast(DomainScalarMap, n1_fpc)
            _n2_deff = cast(DomainScalarMap, n2_deff)
            _n_final = cast(DomainScalarMap, n_final)
            self._size = [
                Size(
                    stratum=str(s),
                    n0=_n0[s],
                    n1_fpc=_n1_fpc[s],
                    n2_deff=_n2_deff[s],
                    n=_n_final[s],
                )
                for s in _n0
            ]

        return self

    def compare_means(
        self,
        mu1: Number,
        mu2: Number,
        *,
        method: OnePropSizeMethod | DomainScalarMap = OnePropSizeMethod.WALD,
        alloc_ratio: Number = 1.0,  # allocation ratio n2/n1
        alpha: Number = 0.05,
        power: Number = 0.80,
    ) -> Self:
        self._param = PopParam.MEAN
        # placeholder for future implementation; kept for API symmetry
        return self
