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)


@pytest.mark.skip(reason=_NOT_IMPLEMENTED)
def test_compare_models_returns_table_sorted_by_first_metric(linear_data):
    from mcm_toolkit.models import compare_models, fit_linear, fit_polynomial

    x, y = linear_data
    X = x.reshape(-1, 1)
    tbl = compare_models(
        {"linear": fit_linear(X, y), "poly2": fit_polynomial(X, y, degree=2)},
        X,
        y,
        cv=0,
    )
    assert set(tbl.index) >= {"linear", "poly2"}
    assert "rmse" in tbl.columns


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_zero_in_y_true():
    """`y_true` 含 0 时 MAPE 无定义——拦住并点名改用 smape，而不是返回 inf。"""
    from mcm_toolkit.models import mape

    with pytest.raises(ValueError, match="不能含 0"):
        mape([0.0, 2.0], [1.0, 2.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))