MCM-ICM_Toolkit/tests/test_models.py

"""models 包测试。"""

from __future__ import annotations

import numpy as np
import pandas as pd
import pytest

_NOT_IMPLEMENTED = "未实现：见 docs/03-roadmap.md L1"


@pytest.mark.skip(reason=_NOT_IMPLEMENTED)
def test_fit_linear_recovers_known_coefficients(linear_data):
    """真值 y = 2x + 1，拟合系数应在 1% 内。"""
    from mcm_toolkit.models import fit_linear

    x, y = linear_data
    r = fit_linear(x.reshape(-1, 1), y)
    assert np.isclose(r.params["coef"][0], 2.0, rtol=0.01)
    assert np.isclose(r.params["intercept"], 1.0, atol=0.05)


@pytest.mark.skip(reason=_NOT_IMPLEMENTED)
def test_fit_polynomial_matches_numpy(linear_data):
    from mcm_toolkit.models import fit_polynomial

    x, y = linear_data
    r = fit_polynomial(x.reshape(-1, 1), y, degree=2)
    assert np.allclose(r.params["coef"], np.polyfit(x, y, 2), rtol=1e-6)


def test_metrics_known_values():
    from mcm_toolkit.models import mae, r2_score, rmse

    y_true = np.array([1.0, 2.0, 3.0])
    y_pred = np.array([1.0, 2.0, 5.0])
    assert np.isclose(rmse(y_true, y_pred), np.sqrt(4 / 3))
    assert np.isclose(mae(y_true, y_pred), 2 / 3)
    assert r2_score(y_true, y_pred) < 1.0


def test_r2_is_one_for_perfect_fit():
    from mcm_toolkit.models import r2_score

    y = np.array([1.0, 2.0, 3.0, 4.0])
    assert np.isclose(r2_score(y, y), 1.0)


class _MeanModel:
    """最笨的模型：永远预测训练集的均值（`compare_models` 的对照组）。"""

    def fit(self, X, y):
        self.mu_ = float(np.mean(y))
        return self

    def predict(self, X):
        return np.full(X.shape[0], self.mu_)


class _LinearModel:
    """一列特征的普通最小二乘（`compare_models` 的"好模型"）。"""

    def fit(self, X, y):
        self.coef_ = np.polyfit(X[:, 0], y, 1)
        return self

    def predict(self, X):
        return np.polyval(self.coef_, X[:, 0])


def test_compare_models_returns_table_sorted_by_first_metric(linear_data):
    from mcm_toolkit.models import compare_models

    x, y = linear_data
    tbl = compare_models(
        {"mean": _MeanModel(), "linear": _LinearModel()}, x.reshape(-1, 1), y, cv=0
    )

    assert list(tbl.columns) == ["rmse", "mae", "r2"]
    assert list(tbl.index) == ["linear", "mean"]  # 首列 rmse 升序 → 好的模型在前
    assert tbl.loc["linear", "rmse"] < tbl.loc["mean", "rmse"]


def test_metrics_reject_mismatched_shapes():
    from mcm_toolkit.models import rmse

    with pytest.raises(ValueError, match="形状必须一致"):
        rmse([1.0, 2.0], [1.0, 2.0, 3.0])


def test_metrics_reject_empty_input():
    from mcm_toolkit.models import rmse

    with pytest.raises(ValueError, match="不能为空"):
        rmse(np.array([]), np.array([]))


def test_metrics_reject_nan_instead_of_dropping_it():
    """含 NaN 一律报错，**绝不静默删点**：删了样本量就变了，指标不可比。"""
    from mcm_toolkit.models import rmse

    with pytest.raises(ValueError, match="NaN"):
        rmse([1.0, np.nan, 3.0], [1.0, 2.0, 3.0])
    with pytest.raises(ValueError, match="NaN"):
        rmse([1.0, 2.0, 3.0], [1.0, np.nan, 3.0])


def test_metrics_reject_non_whitelisted_type():
    """`str` 也是可迭代对象，必须被入口白名单挡住，不能被当成三个字符算下去。"""
    from mcm_toolkit.models import rmse

    with pytest.raises(TypeError):
        rmse("abc", "abc")  # type: ignore[arg-type]


@pytest.mark.parametrize("name", ["rmse", "mae", "mape", "smape", "r2_score"])
def test_metrics_agree_across_input_types(name):
    """`DataFrame` / `Series` / `ndarray` / `list` 四种入参必须给出同一个数。"""
    import mcm_toolkit.models as models

    yt = [1.0, 2.0, 3.0, 4.0]
    yp = [1.1, 1.9, 3.2, 3.8]
    fn = getattr(models, name)
    base = fn(yt, yp)

    assert np.isclose(fn(np.array(yt), np.array(yp)), base)
    assert np.isclose(fn(pd.Series(yt), pd.Series(yp)), base)
    assert np.isclose(fn(pd.DataFrame({"y": yt}), pd.DataFrame({"y": yp})), base)


def test_mape_rejects_non_positive_y_true():
    """`y_true` 非正时 MAPE 无定义 —— 拦住并点名改用 smape，而不是返回负数或 inf。"""
    from mcm_toolkit.models import mape

    with pytest.raises(ValueError, match="非正数"):
        mape([0.0, 2.0], [1.0, 2.0])
    with pytest.raises(ValueError, match="非正数"):
        mape([-1.0, 2.0], [1.0, 2.0])


def test_metrics_reject_infinite_input():
    """`Inf` 与 `NaN` 一样直接拦：算出来是 `inf` 的指标没有任何意义。"""
    from mcm_toolkit.models import rmse

    with pytest.raises(ValueError, match="Inf"):
        rmse([1.0, np.inf], [1.0, 2.0])
    with pytest.raises(ValueError, match="Inf"):
        rmse([1.0, 2.0], [1.0, -np.inf])


def test_metrics_require_at_least_two_samples():
    """1 个点的误差指标没有统计意义；标量也会被这条拦住（不是被当成 1 个样本）。"""
    from mcm_toolkit.models import rmse

    with pytest.raises(ValueError, match="至少需要包含 2 个样本"):
        rmse([1.0], [1.0])
    with pytest.raises(ValueError, match="至少需要包含 2 个样本"):
        rmse(1.0, 1.0)


def test_smape_is_silent_when_denominator_is_zero():
    """分子分母同时为 0 的点记 0，且**不能漏出 RuntimeWarning**。

    本仓库 `pyproject.toml` 里开了 `filterwarnings = ["error::RuntimeWarning"]`，
    所以这条测试同时守着"零除"这个坑：`np.where` 的两条分支都会被求值，
    直接写除式就会把告警漏出来，测试会红。
    """
    from mcm_toolkit.models import smape

    assert smape([0.0, 0.0], [0.0, 0.0]) == 0.0
    assert smape([0.0, 2.0], [0.0, 2.0]) == 0.0
    assert np.isclose(smape([1.0, 1.0], [0.0, 1.0]), 1.0)


def test_adjusted_r2_rejects_invalid_n_features():
    from mcm_toolkit.models import adjusted_r2

    yt = np.array([1.0, 2.0, 3.0, 4.0, 5.0])
    yp = np.array([1.5, 2.5, 2.5, 4.5, 5.5])

    for bad in (True, 1.0, np.float64(1.0)):
        with pytest.raises(TypeError):
            adjusted_r2(yt, yp, bad)  # type: ignore[arg-type]
    with pytest.raises(ValueError, match="正整数"):
        adjusted_r2(yt, yp, 0)
    with pytest.raises(ValueError, match="样本量太小"):
        adjusted_r2(yt, yp, 4)  # n = 5，分母 n - k - 1 = 0


def test_adjusted_r2_accepts_numpy_integer():
    from mcm_toolkit.models import adjusted_r2

    yt = np.array([1.0, 2.0, 3.0, 4.0, 5.0])
    yp = np.array([1.5, 2.5, 2.5, 4.5, 5.5])

    assert adjusted_r2(yt, yp, np.int64(1)) == adjusted_r2(yt, yp, 1)


def test_adjusted_r2_penalises_extra_features():
    """同一份数据、同一个 R²，自变量越多，调整 R² 越小（没用的变量只吃惩罚）。"""
    from mcm_toolkit.models import adjusted_r2

    yt = np.array([1.0, 2.0, 3.0, 4.0, 5.0, 6.0])
    yp = np.array([1.2, 2.1, 3.3, 3.8, 5.1, 6.2])

    assert adjusted_r2(yt, yp, 3) < adjusted_r2(yt, yp, 2) < adjusted_r2(yt, yp, 1)


def test_r2_matches_analytic_formula():
    """`R² = 1 - SS_res / SS_tot`，且 `SS_tot` 的基准必须是 **y_true 的均值**。

    这里刻意用**非完美拟合**的数据：完美拟合（R² = 1）和常量 y（R² = 0）都能把
    "基准量抄错对象"这类 bug 掩盖掉（个人易错点 31）。
    """
    from mcm_toolkit.models import r2_score

    yt = np.array([1.0, 2.0, 3.0, 4.0, 5.0])
    yp = np.array([1.5, 2.5, 2.5, 4.5, 5.5])
    ss_res = np.sum((yt - yp) ** 2)
    ss_tot = np.sum((yt - np.mean(yt)) ** 2)

    assert np.isclose(r2_score(yt, yp), 1 - ss_res / ss_tot)
    assert np.isclose(r2_score(yt, yp), 0.875)


def test_r2_matches_sklearn():
    """拿成熟库当**独立的参照物**：自实现指标最容易"自己算自己"。"""
    sk = pytest.importorskip("sklearn.metrics")

    from mcm_toolkit.models import r2_score

    rng = np.random.default_rng(0)
    yt = rng.normal(size=50)
    yp = yt + rng.normal(scale=0.5, size=50)

    assert np.isclose(r2_score(yt, yp), sk.r2_score(yt, yp))


def test_r2_is_negative_when_worse_than_predicting_the_mean():
    from mcm_toolkit.models import r2_score

    assert np.isclose(r2_score([1.0, 2.0, 3.0], [3.0, 2.0, 1.0]), -3.0)


def test_adjusted_r2_matches_analytic_formula():
    from mcm_toolkit.models import adjusted_r2, r2_score

    yt = np.array([1.0, 2.0, 3.0, 4.0, 5.0, 6.0])
    yp = np.array([1.2, 2.1, 3.3, 3.8, 5.1, 6.2])
    n, k = len(yt), 2
    r2 = r2_score(yt, yp)

    assert np.isclose(adjusted_r2(yt, yp, k), 1 - (1 - r2) * (n - 1) / (n - k - 1))


def test_compare_models_treats_r2_names_as_aliases(linear_data):
    """`"r2"` 与 `"r2_score"` 是同一个指标的别名：传哪个、列名就叫哪个，数值相同。"""
    from mcm_toolkit.models import compare_models

    x, y = linear_data
    X = x.reshape(-1, 1)
    short = compare_models({"m": _LinearModel()}, X, y, metrics=("rmse", "r2"), cv=0)
    full = compare_models({"m": _LinearModel()}, X, y, metrics=("rmse", "r2_score"), cv=0)

    assert list(short.columns) == ["rmse", "r2"]
    assert list(full.columns) == ["rmse", "r2_score"]
    assert np.isclose(short.loc["m", "r2"], full.loc["m", "r2_score"])


def test_compare_models_sorts_descending_when_first_metric_is_r2(linear_data):
    """误差类是**越小越好**、R² 是**越大越好**，排序方向必须跟着首列指标走。"""
    from mcm_toolkit.models import compare_models

    x, y = linear_data
    tbl = compare_models(
        {"mean": _MeanModel(), "linear": _LinearModel()},
        x.reshape(-1, 1),
        y,
        metrics=("r2",),
        cv=0,
    )

    assert list(tbl.index) == ["linear", "mean"]  # r2 降序 → 好的模型仍在前


def test_compare_models_rejects_bad_metric_names(linear_data):
    from mcm_toolkit.models import compare_models

    x, y = linear_data
    X = x.reshape(-1, 1)

    with pytest.raises(ValueError, match="不认识的指标名"):
        compare_models({"m": _MeanModel()}, X, y, metrics=("mse",), cv=0)
    with pytest.raises(ValueError, match="别名"):
        compare_models({"m": _MeanModel()}, X, y, metrics=("r2", "r2_score"), cv=0)
    with pytest.raises(TypeError, match="不要传单个字符串"):
        compare_models({"m": _MeanModel()}, X, y, metrics="rmse", cv=0)  # type: ignore[arg-type]
    with pytest.raises(TypeError, match="tuple / list"):
        compare_models({"m": _MeanModel()}, X, y, metrics=None, cv=0)  # type: ignore[arg-type]


def test_compare_models_rejects_bad_models_and_cv(linear_data):
    from mcm_toolkit.models import compare_models

    x, y = linear_data
    X = x.reshape(-1, 1)

    with pytest.raises(TypeError, match="字典"):
        compare_models([_MeanModel()], X, y)  # type: ignore[arg-type]
    with pytest.raises(ValueError, match="不能为空"):
        compare_models({}, X, y)
    with pytest.raises(TypeError, match="Model 协议"):
        compare_models({"bad": object()}, X, y, cv=0)
    with pytest.raises(ValueError, match="cv=1"):
        compare_models({"m": _MeanModel()}, X, y, cv=1)
    with pytest.raises(ValueError, match="太大"):
        compare_models({"m": _MeanModel()}, X, y, cv=1000)
    with pytest.raises(ValueError, match="样本数必须一致"):
        compare_models({"m": _MeanModel()}, X, y[:-1], cv=0)


def test_compare_models_needs_two_samples_per_fold():
    """每折至少 2 个样本（指标自己的入口要求），所以 cv 最多 n // 2。"""
    from mcm_toolkit.models import compare_models

    X = np.arange(4, dtype=float).reshape(-1, 1)
    y = np.arange(4, dtype=float)

    assert compare_models({"m": _MeanModel()}, X, y, cv=2, random_state=0).shape == (1, 3)
    with pytest.raises(ValueError, match="太大"):
        compare_models({"m": _MeanModel()}, X, y, cv=3, random_state=0)
    with pytest.raises(ValueError, match="样本太少"):
        compare_models({"m": _MeanModel()}, X[:1], y[:1], cv=0)


def test_compare_models_is_reproducible_with_random_state(linear_data):
    """分折用了随机数，所以必须服从 `random_state`（架构 §4.2）。"""
    from mcm_toolkit.models import compare_models

    x, y = linear_data
    X = x.reshape(-1, 1)
    models = {"mean": _MeanModel(), "linear": _LinearModel()}

    first = compare_models(models, X, y, cv=5, random_state=42)
    second = compare_models(models, X, y, cv=5, random_state=42)

    assert first.equals(second)