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)


_VALID_CORR_METHODS = ("pearson", "kendall", "spearman")


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
    ----------
    method:
        `"normal"`：`mean ± z·s/√n`（n 大时用）；
        `"t"`：小样本用 t 分布（更保守，n < 30 建议用这个）；
        `"proportion"`：比例的标准误 `√(p(1-p)/n)`。

    Examples
    --------
    >>> lo, hi = confidence_interval([1, 2, 3, 4, 5], confidence=0.95)
    """
    # TODO(L0)
    raise NotImplementedError


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` 时是单样本检验。

    Returns
    -------
    dict：`{"statistic", "pvalue", "df", "significant", "conclusion"}`。
    `conclusion` 是一句能直接抄进论文的中文结论。
    """
    # TODO(L1)
    raise NotImplementedError


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