MCM-ICM_Toolkit/tests/test_data.py

"""data 包测试。

L0 阶段的约定：**只测已实现的行为**。未实现的算法用逐条 skip 标出，
不要用模块级 pytestmark —— 那会把同文件里已实现的函数一起跳过，
让 load_table 这种已经写好的函数变成零覆盖。
"""

from __future__ import annotations

from pathlib import Path

import numpy as np
import pandas as pd
import pytest

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


# ---------------------------------------------------------------- io: load_table


def test_load_table_detects_gbk(tmp_path):
    """国内赛题数据常见 GBK，编码探测必须能兜住。"""
    from mcm_toolkit.data import load_table

    p = tmp_path / "gbk.csv"
    p.write_bytes("城市,人口\n北京,2189\n".encode("gbk"))
    df = load_table(p)
    assert list(df.columns) == ["城市", "人口"]


def test_load_table_keeps_utf8_readable(tmp_path):
    """UTF-8 文件必须按 UTF-8 读出，不能先试 GBK 变成乱码。"""
    from mcm_toolkit.data import load_table

    p = tmp_path / "utf8.csv"
    p.write_text("省份,GDP增速\n广东,5.8\n", encoding="utf-8")
    df = load_table(p)
    assert list(df.columns) == ["省份", "GDP增速"]
    assert df.loc[0, "省份"] == "广东"


def test_load_table_maps_na_like_tokens_to_nan(tmp_path):
    """赛题里的缺失值长得五花八门，默认 na_values 要全部吃成 NaN。"""
    from mcm_toolkit.data import load_table

    p = tmp_path / "na.csv"
    p.write_text("a,b,c\n1,NA,\n2,-,#N/A\n", encoding="utf-8")
    df = load_table(p)
    # 外面套 bool()：df["b"] 的类型是 Series | DataFrame，直接当条件用 pyright 会报
    # "条件操作数无效"；套上 bool() 既满足类型检查，又能在它真返回 Series 时当场炸掉。
    assert bool(df["b"].isna().all())
    assert bool(df["c"].isna().all())
    assert bool(df["a"].notna().all())


def test_load_table_respects_explicit_encoding(tmp_path):
    """显式给了 encoding 就照办，不再做探测。"""
    from mcm_toolkit.data import load_table

    p = tmp_path / "gbk.csv"
    p.write_bytes("城市,人口\n上海,2487\n".encode("gbk"))
    df = load_table(p, encoding="gbk")
    assert list(df.columns) == ["城市", "人口"]


def test_load_table_rejects_unknown_suffix(tmp_path):
    """不支持的扩展名要抛 ValueError，而且要带上后缀名。"""
    from mcm_toolkit.data import load_table

    p = tmp_path / "data.txt"
    p.write_text("a,b\n1,2\n", encoding="utf-8")
    with pytest.raises(ValueError, match=r"\.txt"):
        load_table(p)


def test_load_table_missing_file_raises(tmp_path):
    """文件不存在时要把路径原样带出来，方便定位。"""
    from mcm_toolkit.data import load_table

    missing = tmp_path / "nope.csv"
    with pytest.raises(FileNotFoundError, match=r"nope\.csv"):
        load_table(missing)


# ---------------------------------------------------------------- io: save_table


def test_save_table_round_trip(tmp_path):
    """存进去再读出来，内容要一致，并且返回实际写入的 Path。"""
    from mcm_toolkit.data import load_table, save_table

    df = pd.DataFrame({"城市": ["北京", "上海"], "人口": [2189.0, 2487.0]})
    out = save_table(df, tmp_path / "out.csv")
    assert isinstance(out, Path)
    assert out.exists()
    back = load_table(out)
    assert list(back.columns) == ["城市", "人口"]
    assert back["人口"].tolist() == [2189.0, 2487.0]


def test_save_table_rejects_non_dataframe(tmp_path):
    """传错类型要当场报错，不要写出一个奇怪的文件。"""
    from mcm_toolkit.data import save_table

    with pytest.raises(TypeError):
        save_table([1, 2, 3], tmp_path / "bad.csv")  # type: ignore[arg-type]  # 就是要传错


def test_save_table_validates_before_touching_disk(tmp_path):
    """校验必须早于副作用：参数不合法时，一个目录、一个文件都不能留下。"""
    from mcm_toolkit.data import save_table

    with pytest.raises(ValueError, match=r"\.txt"):
        save_table(pd.DataFrame({"a": [1]}), tmp_path / "d1" / "x.txt")
    with pytest.raises(TypeError):
        save_table([1, 2, 3], tmp_path / "d2" / "x.csv")  # type: ignore[arg-type]  # 就是要传错

    assert sorted(p.name for p in tmp_path.iterdir()) == []


# --------------------------------------------------------------- data: clean


def test_clean_missing_default_drops_any_missing():
    """不传 threshold 时是"只要有缺失就删"。"""
    from mcm_toolkit.data import clean_missing

    df = pd.DataFrame({"a": [1.0, np.nan, 3.0], "b": [1.0, 2.0, np.nan]})
    out = clean_missing(df)
    assert len(out) == 1
    assert list(out.columns) == ["a", "b"]


def test_clean_missing_threshold_keeps_rows_within_limit():
    """四行三列、阈值 0.5：缺 2/3 的两行要删，剩 2 行。

    3 列时 min_non_na 必须是 ceil(3*0.5)=2；写成 int() 截断就会漏删。
    """
    from mcm_toolkit.data import clean_missing

    df = pd.DataFrame(
        {
            "a": [1.0, np.nan, np.nan, 4.0],
            "b": [1.0, np.nan, np.nan, 4.0],
            "c": [1.0, 2.0, np.nan, 4.0],
        }
    )
    out = clean_missing(df, axis=0, threshold=0.5)
    assert len(out) == 2
    assert out["a"].tolist() == [1.0, 4.0]


def test_clean_missing_axis1_drops_empty_column():
    """axis=1 按列删：整列全空（缺失率 1.0 > 0.5）要掉。"""
    from mcm_toolkit.data import clean_missing

    df = pd.DataFrame({"x": [1.0, 2.0, 3.0], "y": [np.nan, np.nan, np.nan]})
    assert list(clean_missing(df, axis=1, threshold=0.5).columns) == ["x"]


def test_clean_missing_rejects_bad_axis():
    from mcm_toolkit.data import clean_missing

    with pytest.raises(ValueError, match="axis"):
        clean_missing(pd.DataFrame({"a": [1.0]}), axis=2)


def test_clean_missing_rejects_non_dataframe():
    from mcm_toolkit.data import clean_missing

    with pytest.raises(TypeError):
        clean_missing([[1.0, 2.0]])


def test_impute_missing_mean_and_ffill():
    from mcm_toolkit.data import impute_missing

    df = pd.DataFrame({"a": [1.0, np.nan, 3.0]})
    assert impute_missing(df, method="mean")["a"].tolist() == [1.0, 2.0, 3.0]
    assert impute_missing(df, method="ffill")["a"].tolist() == [1.0, 1.0, 3.0]


def test_impute_missing_mode_survives_all_nan_column():
    """整列全 NaN 时 mode() 给不出众数：不能崩，也不能硬填一个值进去。"""
    from mcm_toolkit.data import impute_missing

    df = pd.DataFrame({"a": [np.nan, np.nan], "b": [1.0, np.nan]})
    out = impute_missing(df, method="mode")
    assert bool(out["a"].isna().all())
    assert out["b"].tolist() == [1.0, 1.0]


def test_impute_missing_rejects_unknown_method():
    from mcm_toolkit.data import impute_missing

    with pytest.raises(ValueError, match="bogus"):
        impute_missing(pd.DataFrame({"a": [1.0]}), method="bogus")


def test_impute_missing_rejects_non_dataframe():
    from mcm_toolkit.data import impute_missing

    with pytest.raises(TypeError):
        impute_missing([[1.0, 2.0]])


def test_winsorize_clips_series_to_quantiles():
    from mcm_toolkit.data import winsorize

    s = pd.Series([1.0, 2.0, 3.0, 4.0, 1000.0])
    out = winsorize(s, lower=0.1, upper=0.9)
    assert float(out.iloc[0]) == pytest.approx(1.4)
    assert float(out.iloc[-1]) == pytest.approx(601.6)
    assert float(out.max()) < 1000.0


def test_winsorize_rejects_bad_quantile_order():
    from mcm_toolkit.data import winsorize

    with pytest.raises(ValueError, match="lower"):
        winsorize(pd.Series([1.0, 2.0]), lower=0.9, upper=0.1)


def test_winsorize_rejects_non_frame_with_typeerror():
    """§4.7：类型不对必须抛 TypeError，不能抛 ValueError。"""
    from mcm_toolkit.data import winsorize

    with pytest.raises(TypeError):
        winsorize([1.0, 2.0, 3.0])


def test_detect_outliers_rejects_unknown_method():
    """方法名校验在取默认取值之前，所以这条现在就是绿的。"""
    from mcm_toolkit.data import detect_outliers

    with pytest.raises(ValueError, match="bogus"):
        detect_outliers([1.0, 2.0], method="bogus")  # type: ignore[arg-type]


def test_detect_outliers_rejects_bad_type():
    """§4.7：类型不对必须抛 TypeError，不能让它一路飘进 pandas 里去。"""
    from mcm_toolkit.data import detect_outliers

    for bad in (42, 3.14, "123", None, {"a": [1, 2]}):
        with pytest.raises(TypeError):
            detect_outliers(bad)


def test_detect_outliers_survives_degenerate_columns():
    """全 NaN 列与单元素列都不能抛异常，判定一律为 False。"""
    from mcm_toolkit.data import detect_outliers

    df = pd.DataFrame({"all_nan": [np.nan, np.nan, np.nan], "one": [1.0, 2.0, 3.0]})
    assert not bool(np.asarray(detect_outliers(df)).any())
    assert not bool(np.asarray(detect_outliers([5.0])).any())
    assert not bool(np.asarray(detect_outliers([5.0], method="zscore")).any())


def test_detect_outliers_accepts_tuple_as_1d():
    """元组按一维数组型处理，返回值与传 list 一致。"""
    from mcm_toolkit.data import detect_outliers

    mask = np.asarray(detect_outliers((1.0, 2.0, 3.0, 4.0, 1000.0)))
    assert mask.shape == (5,)
    assert mask.dtype == bool
    assert bool(mask[-1])


# ------------------------ 曾经 xfail 的契约，现已全部实现（留作回归）
#
# 这些用例写的是**文档里已经承诺的行为**。实现没跟上时用 xfail，让它们照样出现在
# 报告里（显示成 x）；修好之后会自动变成 XPASS —— 看到 XPASS 就把 mark 摘掉。
#
# 已摘掉：threshold 范围校验、IQR 的 UnboundLocalError、(n,1) 形状、zscore 默认阈值、
# winsorize 的文本列、detect_outliers 的文本列、impute_missing 的 mode 就地修改。
# **当前 0 条 xfail** —— 以后新增契约照旧：先写测试 + xfail，实现跟上就摘 mark。


def test_detect_outliers_finds_planted_points():
    """人工埋 3 个离群点，IQR 法应正好命中 3 个。"""
    from mcm_toolkit.data import detect_outliers

    x = np.concatenate([np.arange(100.0), [1e6, -1e6, 1e7]])
    mask = np.asarray(detect_outliers(x, method="iqr"))
    assert int(mask.sum()) == 3
    assert bool(mask[-3:].all())


def test_detect_outliers_returns_mask_not_filtered_data():
    """约定：只标不删——返回与原数据同形的布尔掩码。"""
    from mcm_toolkit.data import detect_outliers

    x = np.arange(50.0)
    mask = np.asarray(detect_outliers(x, method="zscore"))
    assert mask.shape == x.shape
    assert mask.dtype == bool


def test_detect_outliers_zscore_flags_planted_point():
    from mcm_toolkit.data import detect_outliers

    x = np.concatenate([np.zeros(99), [1e6]])
    mask = np.asarray(detect_outliers(x, method="zscore"))
    assert int(mask.sum()) == 1
    assert bool(mask[-1].all())


def test_detect_outliers_survives_text_columns():
    """赛题表几乎都带文本列：分位数只在数值列上算，文本列恒 False。"""
    from mcm_toolkit.data import detect_outliers

    df = pd.DataFrame({"city": ["a", "b", "c", "d"], "v": [1.0, 2.0, 3.0, 1000.0]})
    mask = detect_outliers(df)
    assert int(np.asarray(mask).sum()) == 1
    assert not bool(np.asarray(mask["city"]).any())
    assert bool(np.asarray(mask["v"]).any())


def test_detect_outliers_all_text_frame_returns_all_false():
    """整表没有数值列时不报错，返回全 False 掩码（语义：这一列没有离群点）。"""
    from mcm_toolkit.data import detect_outliers

    df = pd.DataFrame({"a": ["x", "y"], "b": ["z", "w"]})
    mask = detect_outliers(df)
    assert mask.shape == df.shape
    assert not bool(np.asarray(mask).any())


def test_winsorize_survives_text_columns():
    """文本列必须原样留着：分位数只在数值列上算。"""
    from mcm_toolkit.data import winsorize

    df = pd.DataFrame({"city": ["a", "b", "c", "d"], "v": [1.0, 2.0, 3.0, 1000.0]})
    out = winsorize(df)
    assert out["city"].tolist() == ["a", "b", "c", "d"]
    assert float(out["v"].iloc[-1]) < 1000.0


def test_clean_missing_rejects_out_of_range_threshold():
    from mcm_toolkit.data import clean_missing

    with pytest.raises(ValueError, match="threshold"):
        clean_missing(pd.DataFrame({"a": [1.0, np.nan]}), threshold=1.5)


def test_impute_missing_mode_does_not_mutate_input():
    """mode 分支也必须返回新表：不能顺手把调用方的表改了。"""
    from mcm_toolkit.data import impute_missing

    df = pd.DataFrame({"a": [1.0, np.nan, 1.0]})
    out = impute_missing(df, method="mode")
    assert out["a"].tolist() == [1.0, 1.0, 1.0]
    assert bool(df["a"].isna().any())


# ------------------------------------------------------ data: transform


def test_normalize_minmax_range(rng):
    """minmax 是**按列**压到 [0, 1]，不是整表压一次。"""
    from mcm_toolkit.data import normalize

    x = rng.normal(size=(50, 3))
    out = np.asarray(normalize(x, method="minmax"))
    assert np.allclose(out.min(axis=0), 0.0)
    assert np.allclose(out.max(axis=0), 1.0)


def test_normalize_minmax_doctest_values():
    """docstring 里那个 Examples 就是规格，不能只有形似。"""
    from mcm_toolkit.data import normalize

    assert normalize([1, 2, 3]).tolist() == [0.0, 0.5, 1.0]


def test_normalize_respects_feature_range():
    from mcm_toolkit.data import normalize

    out = np.asarray(normalize([1.0, 2.0, 3.0], feature_range=(-1.0, 1.0)))
    assert np.allclose(out, [-1.0, 0.0, 1.0])


def test_normalize_constant_column_stays_finite():
    """常量列 max == min，不能出 nan / inf；固定压到 feature_range 的下界。"""
    from mcm_toolkit.data import normalize

    out = np.asarray(normalize([[2.0, 1.0], [2.0, 3.0], [2.0, 5.0]]))
    assert np.allclose(out[:, 0], 0.0)
    assert bool(np.isfinite(out).all())


def test_normalize_vector_method_gives_unit_columns():
    from mcm_toolkit.data import normalize

    out = np.asarray(normalize([[3.0, 0.0], [4.0, 0.0]], method="vector"))
    assert np.allclose(np.linalg.norm(out, axis=0), [1.0, 0.0])


def test_normalize_max_method_scales_by_column_max():
    from mcm_toolkit.data import normalize

    out = np.asarray(normalize([[1.0, 10.0], [2.0, 20.0]], method="max"))
    assert np.allclose(out, [[0.5, 0.5], [1.0, 1.0]])


def test_normalize_rejects_bad_method_and_feature_range():
    from mcm_toolkit.data import normalize

    with pytest.raises(ValueError, match="bogus"):
        normalize([1.0, 2.0], method="bogus")
    with pytest.raises(ValueError, match="feature_range"):
        normalize([1.0, 2.0], feature_range=(1.0, 0.0))
    with pytest.raises(TypeError, match="feature_range"):
        normalize([1.0, 2.0], feature_range=(0.0, 1.0, 2.0))  # type: ignore[arg-type]


def test_standardize_has_zero_mean_unit_std(rng):
    from mcm_toolkit.data import standardize

    x = rng.normal(5, 3, size=1000)
    out = np.asarray(standardize(x))
    assert abs(out.mean()) < 1e-9
    assert abs(out.std() - 1.0) < 1e-9


def test_standardize_is_per_column(rng):
    """按列标准化：每列各自均值 0、标准差 1。"""
    from mcm_toolkit.data import standardize

    x = rng.normal(size=(200, 3)) * [1.0, 10.0, 100.0] + [0.0, 50.0, -3.0]
    out = np.asarray(standardize(x))
    assert np.allclose(out.mean(axis=0), 0.0, atol=1e-12)
    assert np.allclose(out.std(axis=0), 1.0, atol=1e-12)


def test_standardize_honours_ddof():
    """ddof 要真的传下去：ddof=0 才得到单位方差；ddof=1 时结果的 std = sqrt((n-1)/n)。

    除以更大的样本标准差，结果自然更小 —— 所以 ddof=1 时 out.std() **小于** 1。
    """
    from mcm_toolkit.data import standardize

    x = np.array([1.0, 2.0, 3.0, 4.0])
    assert np.allclose(standardize(x, ddof=0).std(), 1.0)
    assert np.allclose(standardize(x, ddof=1).std(), np.sqrt(3.0 / 4.0))


def test_standardize_constant_column_stays_finite():
    from mcm_toolkit.data import standardize

    out = np.asarray(standardize([[2.0, 1.0], [2.0, 3.0]]))
    assert np.allclose(out[:, 0], 0.0)
    assert bool(np.isfinite(out).all())


def test_standardize_rejects_bad_type():
    from mcm_toolkit.data import standardize

    with pytest.raises(TypeError):
        standardize("abc")


def test_log_transform_shifts_to_keep_domain():
    """含 0 或负值时自动取 offset = 1 - min(x)，保证 log 的定义域。"""
    from mcm_toolkit.data import log_transform

    x = np.array([-5.0, 0.0, 5.0])
    out = np.asarray(log_transform(x))
    assert bool(np.isfinite(out).all())
    assert np.allclose(out, np.log(x + 6.0))


def test_log_transform_positive_data_needs_no_offset():
    from mcm_toolkit.data import log_transform

    assert np.allclose(log_transform([1.0, np.e]), [0.0, 1.0])


def test_log_transform_rejects_bad_type():
    from mcm_toolkit.data import log_transform

    with pytest.raises(TypeError):
        log_transform("abc")


def test_one_hot_encode_auto_detects_text_columns_quietly():
    """pandas 3 起字符串列的 dtype 是 'str'：要能选中，且不能走已弃用的 select_dtypes 路径。"""
    import warnings

    from mcm_toolkit.data import one_hot_encode

    df = pd.DataFrame({"city": ["A", "B", "C"], "v": [1.0, 2.0, 3.0]})
    with warnings.catch_warnings(record=True) as caught:
        warnings.simplefilter("always")
        out = one_hot_encode(df)

    messages = [str(w.message) for w in caught]
    assert not messages, f"自动识别不该产生任何告警，实际：{messages}"
    assert "city_A" in out.columns
    assert "v" in out.columns


def test_one_hot_encode_explicit_columns_and_drop_first():
    from mcm_toolkit.data import one_hot_encode

    df = pd.DataFrame({"city": ["A", "B", "C"], "v": [1.0, 2.0, 3.0]})
    out = one_hot_encode(df, columns=["city"], drop_first=True)
    assert "city_A" not in out.columns
    assert "city_B" in out.columns
    assert "city_C" in out.columns


def test_one_hot_encode_rejects_missing_column_and_non_frame():
    from mcm_toolkit.data import one_hot_encode

    df = pd.DataFrame({"city": ["A"], "v": [1.0]})
    with pytest.raises(ValueError, match="nope"):
        one_hot_encode(df, columns=["nope"])
    with pytest.raises(TypeError):
        one_hot_encode([1, 2, 3])


# ------------------------------------------------- 以下为未实现算法（L0/L1 待写）


@pytest.mark.skip(reason=_NOT_IMPLEMENTED)
def test_train_test_split_sizes_and_no_overlap(rng):
    from mcm_toolkit.data import train_test_split

    X = rng.normal(size=(100, 2))
    Xtr, Xte = train_test_split(X, test_size=0.2, random_state=0)
    assert len(Xtr) == 80 and len(Xte) == 20
    assert len({tuple(r) for r in Xtr} & {tuple(r) for r in Xte}) == 0