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",
    "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 correlation_matrix(
    data: Any,
    method: Literal["pearson", "kendall", "spearman"] = "pearson",
    min_periods: int = 1,
) -> Any:
    """相关矩阵 + 显著性星标。

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

    Returns
    -------
    DataFrame，值 = 相关系数；`min_periods` 控制缺失值容差。
    ⚠️ 当前实现返回的是**带星标的字符串表**（`1.000***` 这种），与上面这句契约不一致 ——
    口径待定，见 `STATUS.md`。

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

    Notes
    -----
    星标读法：`***` p < 0.001，`**` p < 0.01，`*` p < 0.05，其余不带星。

    Examples
    --------
    >>> correlation_matrix(df).round(3)       # doctest: +SKIP
    """
    if not isinstance(data, pd.DataFrame):
        raise TypeError(f"data必须是 pd.DataFrame,而你传的是{type(data)}")

    valid_method = ["pearson", "kendall", "spearman"]
    if method not in valid_method:
        raise ValueError(f"不支持的方法:{method},仅支持{valid_method}")

    numeric_cols = data.select_dtypes(include=[np.number]).columns
    if len(numeric_cols) == 0:
        raise ValueError("没有找到数值列")
    # data[Index] 在 pandas 存根里是 DataFrame | Series 重载，按运行时真值收窄（坑 8）
    df_num = cast(pd.DataFrame, data[numeric_cols])

    corr_matrix = df_num.corr(method=method, min_periods=min_periods)

    p_values = pd.DataFrame(0.0, index=corr_matrix.index, columns=corr_matrix.columns, dtype=float)
    for i in range(len(corr_matrix.columns)):
        for j in range(len(corr_matrix.columns)):
            if i == j:
                continue
            if method == "pearson":
                p = _pvalue(stats.pearsonr(df_num.iloc[:, i], df_num.iloc[:, j]))
            elif method == "kendall":
                p = _pvalue(stats.kendalltau(df_num.iloc[:, i], df_num.iloc[:, j]))
            else:
                p = _pvalue(stats.spearmanr(df_num.iloc[:, i], df_num.iloc[:, j]))
            p_values.iloc[i, j] = p

    def add_stars(val: float, p: float) -> str:
        # 保留三位小数
        val_str = f"{val:.3f}"
        if p < 0.001:
            return val_str + "***"
        if p < 0.01:
            return val_str + "**"
        if p < 0.05:
            return val_str + "*"
        return val_str

    formatted_corr = corr_matrix.copy().astype(str)
    for i in range(len(corr_matrix)):
        for j in range(len(corr_matrix)):
            formatted_corr.iloc[i, j] = add_stars(corr_matrix.iloc[i, j], p_values.iloc[i, j])

    return formatted_corr

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