# src/svy/engine/categorical/ranktesting.py
"""
Engine for design-based rank tests under complex sampling.

Implements Lumley & Scott (2013), Biometrika 100(4), 831-842.

The algorithm:
    1. Compute estimated population mid-ranks using survey weights.
    2. Apply a rank-score transformation (Wilcoxon, van der Waerden, median, or custom).
    3. For two groups: fit weighted OLS of rank scores on group indicator,
       compute design-based SE via influence functions and svytotal, return t-test.
    4. For k groups: fit weighted OLS of rank scores on factor dummies,
       compute design-based variance-covariance via sandwich estimator, return Wald/F-test.
"""

from __future__ import annotations

from typing import Callable

import numpy as np

from scipy import stats as sp_stats

from svy.core.enumerations import RankScoreMethod


# =============================================================================
# Score functions
# =============================================================================


def _wilcoxon_score(r: np.ndarray, N: float) -> np.ndarray:
    """Wilcoxon: g(r) = r / N  (proportional ranks)."""
    return r / N


def _vanderwaerden_score(r: np.ndarray, N: float) -> np.ndarray:
    """Van der Waerden: g(r) = Φ⁻¹(r / N)."""
    # Clip to avoid -inf / +inf at boundaries
    u = np.clip(r / N, 1e-10, 1.0 - 1e-10)
    return sp_stats.norm.ppf(u)


def _median_score(r: np.ndarray, N: float) -> np.ndarray:
    """Mood's median: g(r) = I(r > N/2)."""
    return (r > N / 2).astype(np.float64)


SCORE_FUNCTIONS: dict[RankScoreMethod, Callable[[np.ndarray, float], np.ndarray]] = {
    RankScoreMethod.KRUSKAL_WALLIS: _wilcoxon_score,
    RankScoreMethod.VANDER_WAERDEN: _vanderwaerden_score,
    RankScoreMethod.MEDIAN: _median_score,
}


# =============================================================================
# Estimated population mid-ranks
# =============================================================================


def _compute_estimated_ranks(y: np.ndarray, w: np.ndarray) -> np.ndarray:
    """
    Compute estimated population mid-ranks (R-hat).

    For each observation i:
        R_hat_i = sum_j [ w_j * I(y_j < y_i) + 0.5 * w_j * I(y_j == y_i) ]

    This is equivalent to the R code:
        rankhat[ii] <- ave(cumsum(w[ii]) - w[ii]/2, factor(y[ii]))

    Parameters
    ----------
    y : ndarray of shape (n,)
        Response variable.
    w : ndarray of shape (n,)
        Sampling weights (positive).

    Returns
    -------
    rankhat : ndarray of shape (n,)
        Estimated population mid-ranks, in original order of y.
    """
    n = len(y)

    # Sort by y (stable for tie handling)
    ii = np.argsort(y, kind="mergesort")
    y_sorted = y[ii]
    w_sorted = w[ii]

    # cumsum(w) - w/2 gives weighted mid-rank position
    cumw = np.cumsum(w_sorted)
    midrank_sorted = cumw - w_sorted / 2.0

    # Average over ties (same y value gets average mid-rank)
    # This mirrors R's ave(..., factor(y[ii]))
    rankhat_sorted = np.empty(n, dtype=np.float64)
    i = 0
    while i < n:
        j = i + 1
        while j < n and y_sorted[j] == y_sorted[i]:
            j += 1
        avg_midrank = np.mean(midrank_sorted[i:j])
        rankhat_sorted[i:j] = avg_midrank
        i = j

    # Map back to original order
    rankhat = np.empty(n, dtype=np.float64)
    rankhat[ii] = rankhat_sorted
    return rankhat


# =============================================================================
# Two-sample rank test
# =============================================================================


def _ranktest_two_sample(
    y: np.ndarray,
    g: np.ndarray,
    w: np.ndarray,
    stratum: np.ndarray | None,
    psu: np.ndarray,
    ssu: np.ndarray | None,
    fpc: float,
    score_fn: Callable[[np.ndarray, float], np.ndarray],
    alpha: float = 0.05,
    alternative: str = "two-sided",
) -> dict:
    """
    Two-sample design-based rank test.

    Mirrors R's svyranktest for length(unique(g)) == 2.

    Returns a dict with keys: delta, se, t, df, p_value, estimates.
    """
    from svy.engine.estimation.taylor import _taylor_variance

    # 1. Compute estimated population mid-ranks
    N_hat = np.sum(w)
    rankhat = _compute_estimated_ranks(y, w)

    # 2. Apply score transformation
    rankscore = score_fn(rankhat, N_hat)

    # 3. Weighted OLS: rankscore ~ g
    #    This gives the difference in mean rank score between groups.
    #    Using the WLS formulation: X = [1, g_indicator], weights = w
    levels = np.unique(g)
    # Encode group as 0/1 indicator for the second level
    g_binary = (g == levels[1]).astype(np.float64)

    # Design matrix
    n = len(y)
    X = np.column_stack([np.ones(n), g_binary])

    # Weighted least squares: beta = (X'WX)^{-1} X'W y
    XtW = X.T * w[np.newaxis, :]  # (2, n) — avoids forming full diagonal
    XtWX = XtW @ X  # (2, 2)
    XtWy = XtW @ rankscore  # (2,)
    beta = np.linalg.solve(XtWX, XtWy)
    delta = beta[1]  # difference in mean rank score

    # Fitted values and residuals
    fitted = X @ beta
    resid = rankscore - fitted

    # Covariance unscaled: (X'WX)^{-1}
    cov_unscaled = np.linalg.inv(XtWX)

    # 4. Influence functions for design-based variance
    #    R: infn = (X * resid) @ cov_unscaled, then svytotal(infn) weights by w.
    #    svy's _taylor_variance expects pre-weighted scores (like _scores_total),
    #    so we multiply by w here.
    infn = (X * resid[:, np.newaxis] * w[:, np.newaxis]) @ cov_unscaled  # (n, 2)

    # Design-based variance of the total of influence functions
    var_total = _taylor_variance(
        y_score=infn,
        wgt=w,
        stratum=stratum,
        psu=psu,
        ssu=ssu,
        fpc=fpc,
    )
    # var_total is a (2, 2) matrix; SE of delta is sqrt of [1,1] element
    se_delta = np.sqrt(var_total[1, 1])

    # 5. Degrees of freedom: degf(design) - 1, matching R and ttesting.py
    from svy.engine.categorical.ttesting import _get_subpop_df

    df = _get_subpop_df(stratum, psu) - 1
    df = max(df, 1)  # safety floor

    # 6. Test statistic and p-value
    t_stat = delta / se_delta if se_delta > 0 else np.nan

    if alternative == "two-sided":
        p_value = 2.0 * sp_stats.t.sf(np.abs(t_stat), df=df)
    elif alternative == "less":
        p_value = sp_stats.t.cdf(t_stat, df=df)
    elif alternative == "greater":
        p_value = sp_stats.t.sf(t_stat, df=df)
    else:
        p_value = 2.0 * sp_stats.t.sf(np.abs(t_stat), df=df)

    # 7. Per-group mean rank scores and SEs
    #    beta[0] = mean rank score for levels[0] (reference)
    #    beta[0] + beta[1] = mean rank score for levels[1]
    #    SE from diagonal of the sandwich variance: var_total
    group_means = [float(beta[0]), float(beta[0] + beta[1])]
    # Var(beta[0]) = var_total[0,0]; Var(beta[0]+beta[1]) = sum of var_total
    se_g0 = float(np.sqrt(var_total[0, 0]))
    se_g1 = float(np.sqrt(var_total[0, 0] + 2 * var_total[0, 1] + var_total[1, 1]))
    group_ses = [se_g0, se_g1]

    return {
        "delta": float(delta),
        "se": float(se_delta),
        "t": float(t_stat),
        "df": float(df),
        "p_value": float(p_value),
        "levels": (levels[0], levels[1]),
        "group_means": group_means,
        "group_ses": group_ses,
    }


# =============================================================================
# K-sample rank test
# =============================================================================


def _ranktest_k_sample(
    y: np.ndarray,
    g: np.ndarray,
    w: np.ndarray,
    stratum: np.ndarray | None,
    psu: np.ndarray,
    ssu: np.ndarray | None,
    fpc: float,
    score_fn: Callable[[np.ndarray, float], np.ndarray],
    alpha: float = 0.05,
) -> dict:
    """
    K-sample design-based rank test (Kruskal-Wallis style).

    Mirrors R's multiranktest.

    Returns a dict with keys: ndf, ddf, chisq, f_stat, p_value, levels.
    """
    from svy.engine.estimation.taylor import _taylor_variance

    # 1. Compute estimated population mid-ranks
    N_hat = np.sum(w)
    rankhat = _compute_estimated_ranks(y, w)

    # 2. Apply score transformation
    rankscore = score_fn(rankhat, N_hat)

    # 3. Weighted OLS: rankscore ~ factor(g)
    #    Create dummy variables for k-1 groups (reference = first level)
    levels = np.sort(np.unique(g))
    k = len(levels)
    ndf = k - 1

    # Design matrix: intercept + k-1 dummies
    n = len(y)
    X = np.zeros((n, k), dtype=np.float64)
    X[:, 0] = 1.0  # intercept
    for j in range(1, k):
        X[:, j] = (g == levels[j]).astype(np.float64)

    # Weighted least squares
    XtW = X.T * w[np.newaxis, :]
    XtWX = XtW @ X
    XtWy = XtW @ rankscore
    beta = np.linalg.solve(XtWX, XtWy)

    # 4. Design-based variance-covariance of beta
    #    Using the sandwich estimator via influence functions
    fitted = X @ beta
    resid = rankscore - fitted
    cov_unscaled = np.linalg.inv(XtWX)

    # Influence functions: same as R's cov_inffun_glm
    #   U = residuals * weights * X / sampling_weights
    #   h = U @ (X'WX)^{-1}
    #   V = _taylor_variance(h)
    # For OLS with sampling weights, working residuals * model weights = resid * w
    # and dividing by sampling weights cancels, leaving: resid * X @ cov_unscaled
    infn = (X * (resid * w)[:, np.newaxis]) @ cov_unscaled  # (n, k)

    # Design-based variance of the coefficient vector
    V = _taylor_variance(
        y_score=infn,
        wgt=w,
        stratum=stratum,
        psu=psu,
        ssu=ssu,
        fpc=fpc,
    )

    # 5. Wald test on the non-intercept coefficients
    #    chisq = beta[1:]' V[1:,1:]^{-1} beta[1:]
    beta_test = beta[1:]  # (k-1,)
    V_test = V[1:, 1:]  # (k-1, k-1)
    chisq = float(beta_test @ np.linalg.solve(V_test, beta_test))

    # 6. Degrees of freedom: degf(design) - ndf
    #    R: ddf <- degf(design) - ndf, where degf = n_PSU - n_Strata
    from svy.engine.categorical.ttesting import _get_subpop_df

    design_df = _get_subpop_df(stratum, psu)
    ddf = design_df - ndf
    ddf = max(ddf, 1)  # safety floor

    # F-statistic and p-value
    f_stat = chisq / ndf
    p_value = float(sp_stats.f.sf(f_stat, dfn=ndf, dfd=ddf))

    # 7. Per-group mean rank scores and SEs
    #    beta[0] = mean rank score for levels[0] (reference)
    #    beta[0] + beta[j] = mean rank score for levels[j], j >= 1
    #    Var(beta[0]) = V[0,0]
    #    Var(beta[0] + beta[j]) = V[0,0] + 2*V[0,j] + V[j,j]
    group_means = [float(beta[0])]
    group_ses = [float(np.sqrt(V[0, 0]))]
    for j in range(1, k):
        group_means.append(float(beta[0] + beta[j]))
        group_ses.append(float(np.sqrt(V[0, 0] + 2 * V[0, j] + V[j, j])))

    return {
        "ndf": ndf,
        "ddf": float(ddf),
        "chisq": chisq,
        "f_stat": float(f_stat),
        "p_value": p_value,
        "levels": list(levels),
        "group_means": group_means,
        "group_ses": group_ses,
    }
