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

import logging

from typing import TYPE_CHECKING, Callable, Literal, cast

import numpy as np
import polars as pl

from svy.categorical.ranktest import RankTestKSample, RankTestTwoSample
from svy.categorical.table import Table
from svy.categorical.ttest import TTestOneGroup, TTestTwoGroups
from svy.core.constants import (
    _INTERNAL_CONCAT_SUFFIX,
)
from svy.core.enumerations import (
    RankScoreMethod,
    TableType,
    TableUnits,
)
from svy.core.types import (
    Number,
)
from svy.engine.categorical.tabulation import _tabulate
from svy.engine.categorical.ttesting import (
    _one_sample_one_group,
    _ttest_by,
)
from svy.utils.checks import assert_no_missing, drop_missing
from svy.utils.helpers import (
    _scale_weights_for_units,
)


# For type checkers only
if TYPE_CHECKING:
    from svy.core.sample import Sample

log = logging.getLogger(__name__)


# -------------------------------------------
# Categorical data analysis facade on Sample
# -------------------------------------------


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

    def tabulate(
        self,
        rowvar: str,
        colvar: str | None = None,
        *,
        units: TableUnits = TableUnits.PROPORTION,
        count_total: float | int | None = None,
        alpha: float = 0.05,
        drop_nulls: bool = False,
        use_labels: bool | None = None,
    ) -> Table:
        """
        Create a frequency table or cross-tabulation.

        Parameters
        ----------
        rowvar : str
            Variable for table rows.
        colvar : str | None
            Variable for table columns (creates two-way table). Default None.
        units : TableUnits
            Output units: PROPORTION, PERCENT, or COUNT. Default PROPORTION.
        count_total : float | int | None
            Total for count scaling. Default None.
        alpha : float
            Significance level for confidence intervals. Default 0.05.
        drop_nulls : bool
            If True, drop rows with missing values. Default False.
        use_labels : bool | None
            If True, display labels instead of codes. None uses Table default.

        Returns
        -------
        Table
            Frequency table with estimates and statistics.
        """
        _raw = self._sample._data
        local_data: pl.DataFrame = (
            cast(pl.DataFrame, _raw)
            if not isinstance(_raw, pl.LazyFrame)
            else cast(pl.DataFrame, _raw.collect())
        )
        design = self._sample._design

        # required columns
        cols = [rowvar] + ([colvar] if colvar else []) + design.specified_fields()
        cols = self._sample._dedup_preserve_order(cols)

        # PERFORMANCE: Select only necessary columns early to reduce memory footprint
        local_data = local_data.select(cols)

        if drop_nulls:
            valid_data = drop_missing(df=local_data, cols=cols, treat_infinite_as_missing=True)
        else:
            assert_no_missing(df=local_data, subset=cols)
            valid_data = local_data

        # ensure concatenated keys; get arrays
        # This handles the complexity of multi-column strata/psu automatically
        concat_data, arrs = self._sample._design_arrays_from_concat(
            data=valid_data, design=design, by=None
        )
        stratum = arrs["stratum"]
        psu = arrs["psu"]
        ssu = arrs["ssu"]

        # weights (scale BEFORE _tabulate, same as your current behavior)
        wgt = (
            concat_data[design.wgt].to_numpy()
            if design.wgt
            else np.ones(concat_data.height, dtype=float)
        )
        wgt = _scale_weights_for_units(wgt, units=units, count_total=count_total)

        # variables to tabulate
        # PERFORMANCE: Only select the specific columns needed for the engine
        subset_vars = (
            concat_data.select([rowvar, colvar]) if colvar else concat_data.select([rowvar])
        )

        # compute (note: _tabulate expects proportions; we pre-scaled wgt accordingly)
        cell_rows, stats = _tabulate(
            vars=subset_vars.cast(pl.String),
            wgt=wgt,
            stratum=stratum,
            psu=psu,  # type: ignore[arg-type]
            ssu=ssu,
            fpc=self._sample._fpc if hasattr(self._sample, "_fpc") else 1,
            alpha=alpha,
        )

        # levels for display
        # Note: We rely on concat_data unique values which are consistent with the estimation data
        rowvals = concat_data[rowvar].unique().sort().to_list()
        colvals = concat_data[colvar].unique().sort().to_list() if colvar else None

        # Get metadata for label resolution
        metadata = getattr(self._sample, "_metadata", None)

        table = Table(
            type=TableType.ONE_WAY if colvar is None else TableType.TWO_WAY,
            rowvar=rowvar,
            colvar=colvar,
            estimates=cell_rows,
            stats=stats,
            rowvals=rowvals,
            colvals=colvals,
            alpha=alpha,
            metadata=metadata,
        )

        # Apply use_labels preference if specified
        if use_labels is not None:
            table.use_labels = use_labels

        return table

    def ttest_one_group(
        self,
        y: str,
        mean_h0: Number,
        *,
        by: str | None = None,
        alpha: float = 0.05,
        drop_nulls: bool = False,
    ) -> TTestOneGroup:
        _raw = self._sample._data
        local_data: pl.DataFrame = (
            cast(pl.DataFrame, _raw)
            if not isinstance(_raw, pl.LazyFrame)
            else cast(pl.DataFrame, _raw.collect())
        )
        design = self._sample._design

        # required columns
        cols = [y] + ([by] if by else []) + design.specified_fields()
        cols = self._sample._dedup_preserve_order(cols)
        local_data = local_data.select(cols)

        # missing handling
        if drop_nulls:
            valid_data = drop_missing(df=local_data, cols=cols, treat_infinite_as_missing=True)
        else:
            assert_no_missing(df=local_data, subset=cols)
            valid_data = local_data

        # Ensure concatenated keys exist; pull arrays for stratum/psu/ssu (and by, if needed).
        # NOTE: we don't pass `by` to the estimator here because _one_sample_one_group
        # does not accept a domain/by argument. This only normalizes keys and fixes
        # multi-key designs safely.
        _cc, _ = self._sample._create_concatenated_cols_from_lists(
            data=valid_data,
            design=design,
            by=by,  # create by__key too (future-proofing)
            null_token="__Null__",
            suffix=_INTERNAL_CONCAT_SUFFIX,
            categorical=True,
            drop_original=False,
        )
        concat_data: pl.DataFrame = cast(pl.DataFrame, _cc)

        # arrays for estimator (consistent with how other modules do it)
        wgt = (
            concat_data[design.wgt].to_numpy()
            if design.wgt
            else np.ones(concat_data.height, dtype=float)
        )

        # Robustly handle potential missing design columns (e.g. SRS designs)
        stratum = (
            concat_data[f"stratum{_INTERNAL_CONCAT_SUFFIX}"].to_numpy()
            if design.stratum is not None
            else None
        )
        psu = (
            concat_data[f"psu{_INTERNAL_CONCAT_SUFFIX}"].to_numpy()
            if design.psu is not None
            else np.arange(concat_data.height, dtype=int)
        )
        ssu = (
            concat_data[f"ssu{_INTERNAL_CONCAT_SUFFIX}"].to_numpy()
            if design.ssu is not None
            else None
        )

        # run test (no by-grouping at this level)
        one_group_test = _one_sample_one_group(
            y=concat_data[y].to_numpy(),
            y_name=y,
            mean_h0=mean_h0,
            wgt=wgt,
            stratum=stratum,
            psu=psu,
            ssu=ssu,
            fpc=self._sample._fpc if hasattr(self._sample, "_fpc") else 1,
            alpha=alpha,
        )

        return one_group_test

    def ttest(
        self,
        y: str,
        *,
        mean_h0: Number = 0,
        group: str | None = None,
        y_pair: str | None = None,
        by: str | None = None,
        alpha: float = 0.05,
        alternative: str = "two-sided",
        drop_nulls: bool = False,
    ) -> TTestOneGroup | TTestTwoGroups | list[TTestOneGroup] | list[TTestTwoGroups]:
        """
        Perform a design-based t-test.

        Args:
            y: Name of the continuous variable to test.
            mean_h0: Hypothesized mean (one-sample) or difference (two-sample). Default 0.
            group: Grouping variable for two-sample test. If None, performs one-sample test.
            y_pair: Second variable for paired t-test (computes y - y_pair).
            by: Stratification variable for domain estimation. When provided, returns
                a list of test results (one per domain level).
            alpha: Significance level for confidence intervals. Default 0.05.
            alternative: Alternative hypothesis. One of "two-sided", "less", "greater".
                Default "two-sided".
            drop_nulls: If True, drop rows with missing values. Default False.

        Returns:
            TTestOneGroup for one-sample tests, TTestTwoGroups for two-sample tests.
            When `by` is specified, returns a list of test results.
        """
        _raw = self._sample._data
        local_data: pl.DataFrame = (
            cast(pl.DataFrame, _raw)
            if not isinstance(_raw, pl.LazyFrame)
            else cast(pl.DataFrame, _raw.collect())
        )
        design = self._sample._design

        # required columns (include y_pair if present)
        cols = (
            [y]
            + ([y_pair] if y_pair else [])
            + ([group] if group else [])
            + ([by] if by else [])
            + design.specified_fields()
        )
        cols = self._sample._dedup_preserve_order(cols)
        local_data = local_data.select(cols)

        # missing handling
        if drop_nulls:
            valid_data = drop_missing(df=local_data, cols=cols, treat_infinite_as_missing=True)
        else:
            assert_no_missing(df=local_data, subset=cols)
            valid_data = local_data

        # Create concatenated keys once (consistent with other modules)
        _cc2, _ = self._sample._create_concatenated_cols_from_lists(
            data=valid_data,
            design=design,
            by=by,
            null_token="__Null__",
            suffix=_INTERNAL_CONCAT_SUFFIX,
            categorical=True,
            drop_original=False,
        )
        concat_data: pl.DataFrame = cast(pl.DataFrame, _cc2)

        # Arrays for estimator
        wgt = (
            concat_data[design.wgt].to_numpy()
            if design.wgt
            else np.ones(concat_data.height, dtype=float)
        )

        # Robustly handle design arrays from the concatenated result
        stratum_arr = (
            concat_data[f"stratum{_INTERNAL_CONCAT_SUFFIX}"].to_numpy()
            if design.stratum is not None
            else None
        )
        psu_arr = (
            concat_data[f"psu{_INTERNAL_CONCAT_SUFFIX}"].to_numpy()
            if design.psu is not None
            else np.arange(concat_data.height, dtype=int)
        )
        ssu_arr = (
            concat_data[f"ssu{_INTERNAL_CONCAT_SUFFIX}"].to_numpy()
            if design.ssu is not None
            else None
        )
        by_arr = concat_data[f"by{_INTERNAL_CONCAT_SUFFIX}"].to_numpy() if by is not None else None

        # y / y_pair logic
        y_arr = concat_data[y].to_numpy()
        if y_pair:
            # PERFORMANCE: vectorized numpy subtraction is fast
            y2_arr = concat_data[y_pair].to_numpy()
            y_input = y_arr - y2_arr
            y_name = f"svy_{y}_minus_{y_pair}"
        else:
            y_input = y_arr
            y_name = y

        # group (kept as the raw column; grouping var is not a multi-key spec)
        group_arr = concat_data[group].to_numpy() if group else None

        # Run test (supports by-domains inside _ttest_by)
        _res = _ttest_by(
            y=y_input,
            y_name=y_name,
            mean_h0=mean_h0,
            group=group_arr,
            group_name=group,
            wgt=wgt,
            stratum=stratum_arr,
            psu=psu_arr,
            ssu=ssu_arr,
            by=by_arr,
            by_name=by,
            fpc=self._sample._fpc if hasattr(self._sample, "_fpc") else 1,
            alpha=alpha,
            alternative=alternative,
        )

        # _ttest_by returns a list; unwrap single result when no `by`
        if len(_res) == 1:
            return _res[0]

        return _res  # type: ignore[return-value]

    def ranktest(
        self,
        y: str,
        *,
        group: str,
        method: RankScoreMethod | None = None,
        score_fn: Callable[[np.ndarray, float], np.ndarray] | None = None,
        by: str | None = None,
        alpha: float = 0.05,
        alternative: Literal["two-sided", "less", "greater"] = "two-sided",
        drop_nulls: bool = False,
    ) -> RankTestTwoSample | RankTestKSample | list[RankTestTwoSample] | list[RankTestKSample]:
        """
        Perform a design-based rank test.

        Implements Lumley & Scott (2013) methodology. Automatically selects
        two-sample (Wilcoxon) or k-sample (Kruskal-Wallis) form based on
        the number of unique levels in ``group``.

        Parameters
        ----------
        y : str
            Name of the continuous variable to test.
        group : str
            Grouping variable. Two levels → two-sample test; three or more → k-sample.
        method : RankScoreMethod | None
            Built-in score function: WILCOXON, VANDER_WAERDEN, or MEDIAN.
            Exactly one of ``method`` or ``score_fn`` must be provided.
        score_fn : Callable | None
            Custom score function ``f(r, N) -> scores`` where ``r`` is the array
            of estimated population mid-ranks and ``N`` is the estimated population
            size. Overrides ``method`` when provided.
        by : str | None
            Domain variable. When provided, returns a list of results (one per
            domain level). Default None.
        alpha : float
            Significance level for confidence intervals. Default 0.05.
        alternative : str
            Alternative hypothesis for two-sample tests. One of "two-sided",
            "less", "greater". Ignored for k-sample tests. Default "two-sided".
        drop_nulls : bool
            If True, drop rows with missing values. Default False.

        Returns
        -------
        RankTestTwoSample | RankTestKSample
            Single result when ``by`` is None.
        list[RankTestTwoSample] | list[RankTestKSample]
            List of results when ``by`` is provided.

        Raises
        ------
        MethodError
            If neither or both of ``method`` and ``score_fn`` are specified.

        References
        ----------
        Lumley, T., & Scott, A.J. (2013). Two-sample rank tests under complex
        sampling. Biometrika, 100(4), 831-842.

        Examples
        --------
        >>> # Wilcoxon rank-sum test (two groups, auto-labeled "Wilcoxon")
        >>> result = sample.categorical.ranktest(
        ...     y="tot_exp", group="urbrur", method=RankScoreMethod.KRUSKAL_WALLIS
        ... )

        >>> # Kruskal-Wallis test (many groups, auto-labeled "Kruskal-Wallis")
        >>> result = sample.categorical.ranktest(
        ...     y="tot_exp", group="geo1", method=RankScoreMethod.KRUSKAL_WALLIS
        ... )

        >>> # Mood's median test
        >>> result = sample.categorical.ranktest(
        ...     y="tot_exp", group="urbrur", method=RankScoreMethod.MEDIAN
        ... )

        >>> # Custom score function (upper quartile test)
        >>> def upper_quartile(r, N):
        ...     return (r > 0.75 * N).astype(float)
        >>> result = sample.categorical.ranktest(
        ...     y="tot_exp", group="urbrur", score_fn=upper_quartile
        ... )
        """
        from svy.categorical.ranktest import RankTestKSample, RankTestTwoSample
        from svy.categorical.ttest import DiffEst, GroupLevels, TtestEst
        from svy.engine.categorical.ranktesting import (
            SCORE_FUNCTIONS,
            _ranktest_k_sample,
            _ranktest_two_sample,
        )
        from svy.errors import MethodError

        # --- Validate method / score_fn ---
        if method is None and score_fn is None:
            raise MethodError.invalid_choice(
                where="ranktest",
                param="method / score_fn",
                got=None,
                allowed=list(RankScoreMethod),
                hint="Provide either method=RankScoreMethod.KRUSKAL_WALLIS or a custom score_fn.",
            )
        if method is not None and score_fn is not None:
            raise MethodError.not_applicable(
                where="ranktest",
                method="ranktest",
                reason="cannot specify both 'method' and 'score_fn'",
                param="method / score_fn",
                hint="Use one or the other.",
            )

        # Resolve score function (display name resolved later based on group count)
        if method is not None:
            resolved_fn = SCORE_FUNCTIONS[method]
            resolved_method = method
        else:
            resolved_fn = score_fn
            resolved_method = None

        # --- Data preparation (same pattern as ttest) ---
        _raw = self._sample._data
        local_data: pl.DataFrame = (
            cast(pl.DataFrame, _raw)
            if not isinstance(_raw, pl.LazyFrame)
            else cast(pl.DataFrame, _raw.collect())
        )
        design = self._sample._design

        cols = [y] + [group] + ([by] if by else []) + design.specified_fields()
        cols = self._sample._dedup_preserve_order(cols)
        local_data = local_data.select(cols)

        if drop_nulls:
            valid_data = drop_missing(df=local_data, cols=cols, treat_infinite_as_missing=True)
        else:
            assert_no_missing(df=local_data, subset=cols)
            valid_data = local_data

        _cc, _ = self._sample._create_concatenated_cols_from_lists(
            data=valid_data,
            design=design,
            by=by,
            null_token="__Null__",
            suffix=_INTERNAL_CONCAT_SUFFIX,
            categorical=True,
            drop_original=False,
        )
        concat_data: pl.DataFrame = cast(pl.DataFrame, _cc)

        # Extract arrays
        wgt = (
            concat_data[design.wgt].to_numpy()
            if design.wgt
            else np.ones(concat_data.height, dtype=float)
        )
        stratum_arr = (
            concat_data[f"stratum{_INTERNAL_CONCAT_SUFFIX}"].to_numpy()
            if design.stratum is not None
            else None
        )
        psu_arr = (
            concat_data[f"psu{_INTERNAL_CONCAT_SUFFIX}"].to_numpy()
            if design.psu is not None
            else np.arange(concat_data.height, dtype=int)
        )
        ssu_arr = (
            concat_data[f"ssu{_INTERNAL_CONCAT_SUFFIX}"].to_numpy()
            if design.ssu is not None
            else None
        )
        by_arr = concat_data[f"by{_INTERNAL_CONCAT_SUFFIX}"].to_numpy() if by is not None else None

        y_arr = concat_data[y].to_numpy().astype(np.float64)
        group_arr = concat_data[group].to_numpy()

        # --- Run test (with domain support) ---
        if by_arr is not None:
            by_levels = sorted(set(by_arr))
        else:
            by_levels = [None]

        results = []
        for by_level in by_levels:
            if by_level is not None:
                mask = by_arr == by_level
                y_sub = y_arr[mask]
                g_sub = group_arr[mask]
                w_sub = wgt[mask]
                str_sub = stratum_arr[mask] if stratum_arr is not None else None
                psu_sub = psu_arr[mask]
                ssu_sub = ssu_arr[mask] if ssu_arr is not None else None
            else:
                y_sub, g_sub, w_sub = y_arr, group_arr, wgt
                str_sub, psu_sub, ssu_sub = stratum_arr, psu_arr, ssu_arr

            n_groups = len(np.unique(g_sub))
            _fpc_raw = self._sample._fpc if hasattr(self._sample, "_fpc") else 1
            if isinstance(_fpc_raw, dict):
                fpc_val = 1.0
            else:
                fpc_val = float(_fpc_raw) if _fpc_raw is not None else 1.0

            # Resolve display name based on group count
            if resolved_method is not None:
                if resolved_method == RankScoreMethod.KRUSKAL_WALLIS and n_groups == 2:
                    method_name = "Wilcoxon"
                else:
                    method_name = resolved_method.value
            else:
                method_name = getattr(score_fn, "__name__", "custom")

            # resolved_fn is guaranteed non-None by the validation above
            assert resolved_fn is not None

            if n_groups == 2:
                raw = _ranktest_two_sample(
                    y=y_sub,
                    g=g_sub,
                    w=w_sub,
                    stratum=str_sub,
                    psu=psu_sub,
                    ssu=ssu_sub,
                    fpc=fpc_val,
                    score_fn=resolved_fn,
                    alpha=alpha,
                    alternative=alternative,
                )
                # Build result container
                from scipy.stats import t as t_dist

                from svy.core.containers import TDist

                t_crit = float(t_dist.ppf(1 - alpha / 2, df=raw["df"]))

                # Per-group mean rank score estimates
                group_ests = []
                for idx_g, level in enumerate(raw["levels"]):
                    g_mean = raw["group_means"][idx_g]
                    g_se = raw["group_ses"][idx_g]
                    g_cv = abs(g_se / g_mean) if g_mean != 0 else float("nan")
                    group_ests.append(
                        TtestEst(
                            by=by,
                            by_level=by_level,
                            group=group,
                            group_level=level,
                            y=y,
                            y_level=None,
                            est=g_mean,
                            se=g_se,
                            cv=g_cv,
                            lci=g_mean - t_crit * g_se,
                            uci=g_mean + t_crit * g_se,
                        )
                    )

                result = RankTestTwoSample(
                    y=y,
                    groups=GroupLevels(
                        var=group,
                        levels=(raw["levels"][0], raw["levels"][1]),
                    ),
                    method_name=method_name,
                    alternative=alternative,
                    diff=[
                        DiffEst(
                            y=y,
                            diff=raw["delta"],
                            se=raw["se"],
                            lci=raw["delta"] - t_crit * raw["se"],
                            uci=raw["delta"] + t_crit * raw["se"],
                            by=by,
                            by_level=by_level,
                        )
                    ],
                    estimates=group_ests,
                    stats=TDist(
                        df=raw["df"],
                        value=raw["t"],
                        p_value=raw["p_value"],
                    ),
                    alpha=alpha,
                )
                results.append(result)

            elif n_groups > 2:
                raw = _ranktest_k_sample(
                    y=y_sub,
                    g=g_sub,
                    w=w_sub,
                    stratum=str_sub,
                    psu=psu_sub,
                    ssu=ssu_sub,
                    fpc=fpc_val,
                    score_fn=resolved_fn,
                    alpha=alpha,
                )
                # Use t-based CI for individual group means (consistent with two-sample)
                from scipy.stats import t as t_dist_k

                from svy.core.containers import FDist

                design_df_k = raw["ddf"] + raw["ndf"]  # degf(design)
                t_crit_k = float(t_dist_k.ppf(1 - alpha / 2, df=design_df_k))

                k_group_ests = []
                for idx_g, level in enumerate(raw["levels"]):
                    g_mean = raw["group_means"][idx_g]
                    g_se = raw["group_ses"][idx_g]
                    g_cv = abs(g_se / g_mean) if g_mean != 0 else float("nan")
                    k_group_ests.append(
                        TtestEst(
                            by=by,
                            by_level=by_level,
                            group=group,
                            group_level=level,
                            y=y,
                            y_level=None,
                            est=g_mean,
                            se=g_se,
                            cv=g_cv,
                            lci=g_mean - t_crit_k * g_se,
                            uci=g_mean + t_crit_k * g_se,
                        )
                    )

                result = RankTestKSample(
                    y=y,
                    group_var=group,
                    group_levels=raw["levels"],
                    method_name=method_name,
                    estimates=k_group_ests,
                    stats=FDist(
                        df_num=raw["ndf"],
                        df_den=raw["ddf"],
                        value=raw["f_stat"],
                        p_value=raw["p_value"],
                    ),
                    alpha=alpha,
                )
                results.append(result)

            else:
                raise MethodError.not_applicable(
                    where="ranktest",
                    method="ranktest",
                    reason=f"group variable '{group}' has fewer than 2 levels in this domain",
                    param="group",
                    hint="Ensure the group variable has at least 2 distinct values.",
                )

        if len(results) == 1:
            return results[0]
        return results
