MCM-ICM_Toolkit/src/mcm_toolkit/stats/inference.py

"""统计推断：区间估计、假设检验、分布拟合。

原建议里的 `mean_test`/`proportion_test`/`distribution_pdf`/`distribution_cdf`/
`distribution_ppf`/`sample_distribution` 都是 scipy 的一行转发，已删。
需要 pdf/cdf/ppf 时，用 `fit_distribution(...).dist.pdf(...)` 即可。
"""

from __future__ import annotations

from typing import Any, Literal, cast

import numpy as np
import pandas as pd
import scipy.stats as stats

__all__ = [
    "anova",
    "chi_square_test",
    "confidence_interval",
    "correlation_matrix",
    "correlation_pvalues",
    "correlation_table",
    "fit_distribution",
    "goodness_of_fit",
    "t_test",
]


def _pvalue(result: Any) -> float:
    """把 scipy 检验结果取成 p 值（float）。

    scipy 的 `SignificanceResult` / `PearsonRResult` 是 `_make_tuple_bunch` 在**运行时**
    拼出来的类，类型检查器看不到 `pvalue` 字段 —— 这里显式收成 `Any` 再取，
    比给每一处调用挂 `# type: ignore` 干净，后面 `t_test` / `anova` 也能复用。
    """
    return float(result.pvalue)


def _statistic(result: Any) -> float:
    """把 scipy 检验结果取成**统计量**（float）。

    与 `_pvalue` 同一个原因：`TtestResult` 也是 `_make_tuple_bunch` 运行时拼的类，
    存根里看不到 `statistic` 字段（见 `docs/04-environment.md` 坑 10）。
    """
    return float(result.statistic)


_VALID_CORR_METHODS = ("pearson", "kendall", "spearman")
_VALID_CONFIDENCE_METHODS = ("normal", "t", "proportion")


def _check_corr_args(data: Any, method: str) -> pd.DataFrame:
    """三个 `correlation_*` 共用的入参守卫，返回收窄后的**数值子表**。"""
    if not isinstance(data, pd.DataFrame):
        raise TypeError(f"data必须是 pd.DataFrame,而你传的是{type(data)}")
    if method not in _VALID_CORR_METHODS:
        raise ValueError(f"不支持的方法:{method},仅支持{list(_VALID_CORR_METHODS)}")
    numeric_cols = data.select_dtypes(include=[np.number]).columns
    if len(numeric_cols) == 0:
        raise ValueError("没有找到数值列")
    # data[Index] 在 pandas 存根里是 DataFrame | Series 重载，按运行时真值收窄（坑 8）
    return cast(pd.DataFrame, data[numeric_cols])


def correlation_matrix(
    data: Any,
    method: Literal["pearson", "kendall", "spearman"] = "pearson",
    min_periods: int = 1,
) -> pd.DataFrame:
    """相关系数矩阵（**数值**，不带星标）。

    Parameters
    ----------
    data:
        必须是 `pandas.DataFrame`。**只有数值列参与**，非数值列自动跳过。
    method:
        `"pearson"` / `"kendall"` / `"spearman"`。
    min_periods:
        每对变量至少要有这么多个**共同非缺失**的观测才给相关系数，透传给 `DataFrame.corr`。
        （pandas 对每对变量都是**成对删缺失**，不是整行删。）

    Returns
    -------
    pandas.DataFrame
        对称矩阵，对角线恒为 `1.0`；观测不足的格子是 NaN。

    Raises
    ------
    TypeError
        `data` 不是 `pandas.DataFrame`。
    ValueError
        `data` 里没有数值列，或 `method` 不在支持列表内。

    Notes
    -----
    要"系数 + 显著性星标"的论文表格用 `correlation_table`；要 p 值矩阵用 `correlation_pvalues`。
    之所以拆开：这张表是**数值**，能挑高相关对、喂聚类 / PCA、画热力图；星标字符串表这些都做不了。

    Examples
    --------
    >>> correlation_matrix(df).round(3)       # doctest: +SKIP
    """
    df_num = _check_corr_args(data, method)
    return df_num.corr(method=method, min_periods=min_periods)


def correlation_pvalues(
    data: Any,
    method: Literal["pearson", "kendall", "spearman"] = "pearson",
    min_periods: int = 1,
) -> pd.DataFrame:
    """`correlation_matrix` 对应的 **p 值矩阵**（数值，不含星标）。

    Parameters
    ----------
    data / method / min_periods:
        含义与 `correlation_matrix` 完全一致，两者一定同形。

    Returns
    -------
    pandas.DataFrame
        对称矩阵。**对角线是 NaN** —— "自己和自己相关"没有 p 值可言，
        写成 `0.0` 会被下游读成"极显著"。

    Raises
    ------
    TypeError / ValueError:
        同 `correlation_matrix`。

    Notes
    -----
    每对变量都先**成对删缺失**再算，与 `DataFrame.corr` 的口径一致 ——
    否则 r 和 p 会基于不同的样本量，对不上；成对有效观测少于 `min_periods` 时给 NaN。
    只算上三角再镜像，成对变量的 p 天然对称，省掉一半重复计算。
    """
    df_num = _check_corr_args(data, method)
    corr = df_num.corr(method=method, min_periods=min_periods)

    p_values = pd.DataFrame(np.nan, index=corr.index, columns=corr.columns, dtype=float)
    for i in range(len(corr.columns)):
        for j in range(i + 1, len(corr.columns)):
            pair = df_num.iloc[:, [i, j]].dropna()
            if len(pair) < min_periods:
                continue
            x = pair.iloc[:, 0]
            y = pair.iloc[:, 1]
            if method == "pearson":
                p = _pvalue(stats.pearsonr(x, y))
            elif method == "kendall":
                p = _pvalue(stats.kendalltau(x, y))
            else:
                p = _pvalue(stats.spearmanr(x, y))
            p_values.iloc[i, j] = p
            p_values.iloc[j, i] = p
    return p_values


def _star(value: float, p: float) -> str:
    """把系数与 p 值拼成 `"0.812**"`。p 为 NaN（对角线 / 观测不足）时不加星标。"""
    text = f"{value:.3f}"
    if p < 0.001:
        return text + "***"
    if p < 0.01:
        return text + "**"
    if p < 0.05:
        return text + "*"
    return text


def correlation_table(
    data: Any,
    method: Literal["pearson", "kendall", "spearman"] = "pearson",
    min_periods: int = 1,
) -> pd.DataFrame:
    """可直接贴论文的相关性表：**系数 + 显著性星标**。

    Parameters
    ----------
    data / method / min_periods:
        含义与 `correlation_matrix` 完全一致。

    Returns
    -------
    pandas.DataFrame
        每个格子是字符串，形如 `"0.812**"`。
        星标：`***` p < 0.001、`**` p < 0.01、`*` p < 0.05，其余不带星。
        **对角线只显示 `1.000` 不带星标** —— 自己和自己没有显著性可言。

    Raises
    ------
    TypeError / ValueError:
        同 `correlation_matrix`。

    Notes
    -----
    这张表**只能看**：格子里是字符串，做不了后续计算。
    要数值请用 `correlation_matrix`（系数）与 `correlation_pvalues`（p 值）。

    Examples
    --------
    >>> print(correlation_table(df))           # doctest: +SKIP
    """
    corr = correlation_matrix(data, method=method, min_periods=min_periods)
    p_values = correlation_pvalues(data, method=method, min_periods=min_periods)

    table = corr.astype(str)
    for i in range(len(corr)):
        for j in range(len(corr.columns)):
            table.iloc[i, j] = _star(corr.iloc[i, j], p_values.iloc[i, j])
    return table


def confidence_interval(
    x: Any,
    confidence: float = 0.95,
    method: str = "normal",
) -> tuple[float, float]:
    """均值或比例的置信区间。

    Parameters
    ----------
    x:
        样本。`list` / `tuple` / `numpy.ndarray` / `pandas.Series` / `pandas.DataFrame`；
        传 DataFrame 时**所有数值列会被摊平成一个样本**，非数值列跳过。
    confidence:
        置信水平，取 `(0, 1)` 开区间，默认 0.95。
    method:
        `"normal"`：`mean ± z·s/√n`（n 大时用）；
        `"t"`：小样本用 t 分布（更保守，n < 30 建议用这个）；
        `"proportion"`：比例的标准误 `√(p(1-p)/n)`，结果会被夹到 `[0, 1]`。

    Returns
    -------
    tuple[float, float]
        `(下限, 上限)`，双侧区间。

    Raises
    ------
    TypeError
        `x` 不是受支持的容器类型，或里面的元素转不成数值。
    ValueError
        `x` 里没有可用数值；`confidence` 不在 `(0, 1)` 内；
        `method` 不在支持列表内；样本量不足 2（估不出标准差）。

    Notes
    -----
    **只有 `proportion` 会把结果夹到 `[0, 1]`** —— 比例的区间不可能越界。
    `normal` / `t` 是**均值**的区间，量纲上完全可以落在 `[0, 1]` 之外
    （平均身高 170cm 的区间下界不可能是 0）；夹一下只会给出**错的下界**。

    Examples
    --------
    >>> lo, hi = confidence_interval([1, 2, 3, 4, 5], confidence=0.95)
    """
    if not isinstance(x, (list, tuple, np.ndarray, pd.Series, pd.DataFrame)):
        raise TypeError(f"x 必须是 list / tuple / ndarray / Series / DataFrame,而你传入了{type(x)}")
    if method not in _VALID_CONFIDENCE_METHODS:
        raise ValueError(f"不支持的方法:{method},仅支持{list(_VALID_CONFIDENCE_METHODS)}")
    if not 0 < confidence < 1:
        raise ValueError(f"confidence 必须在 (0, 1) 开区间内,而你传入了{confidence}")

    if isinstance(x, pd.DataFrame):
        numeric_cols = x.select_dtypes(include=[np.number]).columns
        if len(numeric_cols) == 0:
            raise ValueError("x 里没有数值列")
        arr = cast(pd.DataFrame, x[numeric_cols]).to_numpy(dtype=float).ravel()
    else:
        try:
            arr = np.asarray(x, dtype=float).ravel()
        except (ValueError, TypeError) as e:
            raise TypeError(f"x 必须是数值型数据,无法转成 float 数组:{e}") from None

    arr = arr[~np.isnan(arr)]
    n = len(arr)
    if n < 2:
        raise ValueError(f"样本量至少要有 2 个才估得出标准差,而你只有{n}个")

    if method == "proportion":
        p = float(np.mean(arr))
        se = float(np.sqrt(p * (1 - p) / n))
        z = float(stats.norm.ppf(1 - (1 - confidence) / 2))
        return (max(0.0, p - z * se), min(1.0, p + z * se))

    mean = float(np.mean(arr))
    se = float(np.std(arr, ddof=1)) / np.sqrt(n)
    if method == "normal":
        half = float(stats.norm.ppf(1 - (1 - confidence) / 2)) * se
    else:
        half = float(stats.t.ppf(1 - (1 - confidence) / 2, df=n - 1)) * se
    return (mean - half, mean + half)

def t_test(a: Any, b: Any | None = None, value: float = 0.0, equal_var: bool = False, paired: bool = False, alpha: float = 0.05) -> Any:
    """t 检验：两组均值有没有统计学差异。`b=None` 时是单样本检验。

    Parameters
    ----------
    a:
        第一组样本（`DataFrame` / `Series` / list / tuple / `ndarray`），缺失值自动剔除。
    b:
        第二组样本；`None` 表示单样本检验（对照 `value`）。
    value:
        单样本检验的原假设均值，`b is None` 时生效。
    equal_var:
        `True` 走学生 t（假设方差齐性），`False` 走 Welch（不假设，默认，更稳）。
    paired:
        配对样本检验；要求 `a` 与 `b` 等长。
    alpha:
        显著性水平，取 `(0, 1)` 开区间。

    Returns
    -------
    dict
        `{"statistic", "pvalue", "df", "significant", "conclusion"}`。
        `conclusion` 是一句能直接抄进论文的中文结论。

    Raises
    ------
    TypeError
        样本转不成数值数组。
    ValueError
        `alpha` 不在 `(0, 1)` 内；样本里没有数值列；样本量不足 2；
        配对检验时两组长度不一致。

    Notes
    -----
    传 `DataFrame` 时**所有数值列会被摊平成一个样本** —— 多列不同量纲时请先自己取列。
    """
    if not (0 < alpha < 1):
        raise ValueError(f"显著性水平 alpha 必须在 (0, 1) 之间，你传入了 {alpha}")
    def _clean_data(data: Any) -> np.ndarray:
        if isinstance(data, pd.DataFrame):
            numeric_cols = data.select_dtypes(include=[np.number]).columns
            if len(numeric_cols) == 0:
                raise ValueError("没有找到数值列")
            arr = cast(pd.DataFrame, data[numeric_cols]).to_numpy(dtype=float).ravel()
        else:
            try:
                arr = np.asarray(data, dtype=float).ravel()
            except (ValueError, TypeError) as e:
                raise TypeError(f"样本必须是数值型数据,无法转成 float 数组:{e}") from None
        return arr[~np.isnan(arr)]

    a_arr = _clean_data(a)

    if len(a_arr) < 2:
        raise ValueError("样本 a 的数据量少于 2 个，无法进行 t 检验")
    if b is None:
        res = stats.ttest_1samp(a_arr, popmean=value)
        df = len(a_arr) - 1
        test_name = f"单样本 t 检验 (检验均值是否等于 {value})"
    else:
        b_arr = _clean_data(b)
        if paired:
            if len(a_arr) != len(b_arr):
                raise ValueError("配对 t 检验要求两组数据长度一致")
            res = stats.ttest_rel(a_arr, b_arr)
            df = len(a_arr) - 1
            test_name = "配对样本 t 检验"
        else:
            res = stats.ttest_ind(a_arr, b_arr, equal_var=equal_var)
            if equal_var:
                df = len(a_arr) + len(b_arr) - 2
                test_name = "独立双样本 t 检验 (假设方差齐性)"
            else:
                v1, v2 = np.var(a_arr, ddof=1), np.var(b_arr, ddof=1)
                n1, n2 = len(a_arr), len(b_arr)
                df = (v1/n1 + v2/n2)**2 / ((v1/n1)**2/(n1-1) + (v2/n2)**2/(n2-1))
                test_name = "Welch's 独立双样本 t 检验"

    statistic = _statistic(res)
    pvalue = _pvalue(res)
    significant = bool(pvalue < alpha)

    if significant:
        conclusion = f"在 {alpha} 显著性水平下，{test_name} 结果显著 (p={pvalue:.4f} < {alpha})，拒绝原假设，认为两组数据存在统计学上的显著差异。"
    else:
        conclusion = f"在 {alpha} 显著性水平下，{test_name} 结果不显著 (p={pvalue:.4f} >= {alpha})，不能拒绝原假设，认为两组数据没有统计学上的显著差异。"

    # 6. 返回字典
    return {
        "statistic": statistic,
        "pvalue": pvalue,
        "df": float(df),
        "significant": significant,
        "conclusion": conclusion
    }


def anova(*groups: Any, alpha: float = 0.05) -> Any:
    """单因素方差分析：多组均值是否相同。

    Returns
    -------
    dict：`{"F", "pvalue", "conclusion"}`。
    """
    # TODO(L1)
    raise NotImplementedError


def chi_square_test(observed: Any, expected: Any | None = None, alpha: float = 0.05) -> Any:
    """卡方检验：观察频数与期望频数是否一致（分类变量的独立性 / 拟合优度）。"""
    # TODO(L1)
    raise NotImplementedError


def fit_distribution(x: Any, candidates: tuple[str, ...] = ("norm", "lognorm", "gamma", "expon", "uniform")) -> Any:
    """对候选分布逐个拟合，按 AIC（或 KS 统计量）排序。

    Returns
    -------
    `Result`：`data` 是排好序的候选表；`params["best"]` 是最优分布名，
    `extras["dist"]` 是可直接调 `.pdf / .cdf / .ppf / `.rvs` 的 scipy 分布对象。

    Notes
    -----
    这是"反正态题"的标准做法：先看图 → 候选分布拟合 → 比较 AIC → 做 KS 检验。

    Examples
    --------
    >>> r = fit_distribution(data)            # doctest: +SKIP
    >>> r.params["best"]                      # doctest: +SKIP
    'lognorm'
    """
    # TODO(L1)
    raise NotImplementedError


def goodness_of_fit(x: Any, dist: str = "norm", method: str = "ks") -> Any:
    """拟合优度检验：`"ks"`（连续）/ `"chi2"`（分箱）。"""
    # TODO(L1)
    raise NotImplementedError