MCM-ICM_Toolkit/tests/test_solve.py

"""solve 包测试。"""

from __future__ import annotations

import numpy as np
import pytest

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


def _normal_batch(rng, size, **kwargs):
    """`monte_carlo` 的 sampler 契约：收下实现传来的 `size`，一次抽够。"""
    return rng.normal(0, 1, size)


def test_monte_carlo_integral_converges_to_analytic_value():
    """∫₀¹ x² dx = 1/3。n=2e5 时误差应 < 0.01。"""
    from mcm_toolkit.solve import monte_carlo_integral

    r = monte_carlo_integral(lambda X: X[:, 0] ** 2, bounds=[(0, 1)], n=200_000, random_state=0)
    assert r.success
    assert abs(r.metrics["mean"] - 1 / 3) < 0.01   # 体积必须是正的
    assert len(r.data) == 200_000                 # data 是 n 个样本各自的贡献


def test_monte_carlo_integral_odd_dimension_sign():
    """∫∫∫ x dV over [0,1]³ = 1/2 —— 奇数维最容易暴露"体积符号写反"。"""
    from mcm_toolkit.solve import monte_carlo_integral

    r = monte_carlo_integral(lambda X: X[:, 0], bounds=[(0, 1)] * 3, n=200_000, random_state=0)
    assert abs(r.metrics["mean"] - 0.5) < 0.01


def test_monte_carlo_integral_signed_bounds():
    """∫₀² x dx = 2 —— 边界带宽度时，体积算错和符号写反都会现形。"""
    from mcm_toolkit.solve import monte_carlo_integral

    r = monte_carlo_integral(lambda X: X[:, 0], bounds=[(0, 2)], n=200_000, random_state=0)
    assert abs(r.metrics["mean"] - 2.0) < 0.02


def test_monte_carlo_integral_2d_and_params():
    from mcm_toolkit.solve import monte_carlo_integral

    bounds = [(0, 1), (0, 1)]
    r = monte_carlo_integral(lambda X: X[:, 0] * X[:, 1], bounds=bounds, n=200_000, random_state=0)
    assert abs(r.metrics["mean"] - 0.25) < 0.01
    assert r.name == "monte_carlo_integral"
    assert r.params["dimension"] == 2
    assert r.params["bounds"] == bounds


def test_monte_carlo_integral_rejects_bad_bounds():
    """四类非法 bounds 都要在入口抛 ValueError：区间反了 / 漏括号 / 空 / 参差不齐。"""
    from mcm_toolkit.solve import monte_carlo_integral

    for bad in ([(1, 0)], [0, 1], [], [(0, 1), (0, 2, 3)]):
        with pytest.raises(ValueError):
            monte_carlo_integral(lambda X: X[:, 0], bounds=bad, n=100, random_state=0)


def test_monte_carlo_integral_failure_message_stays_failure():
    """失败时 message 不能被包装层改写成"成功"（易错点 22）。"""
    from mcm_toolkit.solve import monte_carlo_integral

    r = monte_carlo_integral(lambda X: 1.0, bounds=[(0, 1)], n=100, random_state=0)
    assert r.success is False
    assert "成功" not in r.message
    assert "失败" in r.message
    assert r.params["dimension"] == 1          # 失败也要留下追溯信息


def test_monte_carlo_integral_accepts_numpy_integer_n():
    """§4.7：`np.int64` 不是 `int` 的子类，注解与运行期都要收（易错点 1）。"""
    from mcm_toolkit.solve import monte_carlo_integral

    r = monte_carlo_integral(lambda X: X[:, 0], bounds=[(0, 1)], n=np.int64(1000), random_state=0)
    assert r.success
    assert len(r.data) == 1000


def test_monte_carlo_probability_of_standard_normal():
    """P(Z > 0) = 0.5，n=5000 时误差应 < 0.05。"""
    from mcm_toolkit.solve import monte_carlo_probability

    r = monte_carlo_probability(lambda z: z > 0, _normal_batch, n=5000, random_state=0)
    assert r.success
    assert abs(r.metrics["probability"] - 0.5) < 0.05


def test_monte_carlo_is_reproducible():
    from mcm_toolkit.solve import monte_carlo_probability

    a = monte_carlo_probability(lambda z: z > 0, _normal_batch, n=1000, random_state=9)
    b = monte_carlo_probability(lambda z: z > 0, _normal_batch, n=1000, random_state=9)
    assert a.metrics["probability"] == b.metrics["probability"]


def test_monte_carlo_returns_n_results_and_documented_metrics():
    from mcm_toolkit.solve import monte_carlo

    r = monte_carlo(_normal_batch, lambda x: x**2, n=2000, random_state=0)
    assert r.success
    assert r.data.shape == (2000,)
    assert set(r.metrics) == {"mean", "std", "median", "p05", "p95", "ci_low", "ci_high"}


def test_monte_carlo_calls_simulator_only_once():
    """`simulator` 只被调用一次 —— 这是"必须向量化"的可执行版本。"""
    from mcm_toolkit.solve import monte_carlo

    counter = {"calls": 0}

    def simulator(params):
        counter["calls"] += 1
        return params

    r = monte_carlo(_normal_batch, simulator, n=1000, random_state=0)
    assert r.success
    assert counter["calls"] == 1


def test_monte_carlo_data_is_reproducible_with_same_seed():
    from mcm_toolkit.solve import monte_carlo

    a = monte_carlo(_normal_batch, lambda x: x, n=2000, random_state=11)
    b = monte_carlo(_normal_batch, lambda x: x, n=2000, random_state=11)
    assert np.array_equal(a.data, b.data)


def test_monte_carlo_follows_set_seed_when_random_state_omitted():
    from mcm_toolkit.core import set_seed
    from mcm_toolkit.solve import monte_carlo

    set_seed(5)
    a = monte_carlo(_normal_batch, lambda x: x, n=500)
    set_seed(5)
    b = monte_carlo(_normal_batch, lambda x: x, n=500)
    assert np.array_equal(a.data, b.data)


def test_monte_carlo_estimates_known_expectation():
    """E[X²] = 1（X ~ N(0, 1)）；n = 2e5 时误差应 < 0.02。"""
    from mcm_toolkit.solve import monte_carlo

    r = monte_carlo(_normal_batch, lambda x: x**2, n=200_000, random_state=0)
    assert abs(r.metrics["mean"] - 1.0) < 0.02


def test_monte_carlo_docstring_example_runs():
    """docstring 的 Examples 必须真的能跑通（此前那一版是坏的）。"""
    from mcm_toolkit.solve import monte_carlo

    draw = lambda rng, size, **k: rng.normal(0, 1, size)  # noqa: E731
    r = monte_carlo(draw, lambda x: x**2, n=10000, random_state=0)
    assert r.success
    assert abs(r.metrics["mean"] - 1.0) < 0.1


def test_monte_carlo_accepts_numpy_integer_n():
    """§4.7：`np.int64` 不是 `int` 的子类，不能漏。"""
    from mcm_toolkit.solve import monte_carlo

    r = monte_carlo(_normal_batch, lambda x: x, n=np.int64(100), random_state=0)
    assert r.success
    assert r.data.shape == (100,)


def test_monte_carlo_counts_boolean_outcomes():
    """布尔结果按 0/1 统计 —— 事件概率是 MC 最常见的用法。"""
    from mcm_toolkit.solve import monte_carlo

    r = monte_carlo(
        lambda rng, size, **k: rng.normal(0, 1, size) > 0,
        lambda x: x,
        n=5000,
        random_state=0,
    )
    assert r.success
    assert r.data.dtype == float
    assert abs(r.metrics["mean"] - 0.5) < 0.05


def test_monte_carlo_rejects_nan_results():
    """结果里出现 NaN 就判失败，不静默给出一列 nan 指标。"""
    from mcm_toolkit.solve import monte_carlo

    r = monte_carlo(
        lambda rng, size, **k: np.r_[np.zeros(size - 1), np.nan],
        lambda x: x,
        n=10,
        random_state=0,
    )
    assert not r.success
    assert "NaN" in r.message


def test_monte_carlo_rejects_bad_arguments():
    from mcm_toolkit.solve import monte_carlo

    with pytest.raises(TypeError):
        monte_carlo(123, lambda x: x)  # type: ignore[arg-type]  # 就是要传错
    with pytest.raises(TypeError):
        monte_carlo(_normal_batch, 123)  # type: ignore[arg-type]  # 就是要传错
    with pytest.raises(TypeError):
        monte_carlo(_normal_batch, lambda x: x, n=True)
    with pytest.raises(TypeError):
        monte_carlo(_normal_batch, lambda x: x, n=10.0)  # type: ignore[arg-type]  # 就是要传错
    with pytest.raises(ValueError):
        monte_carlo(_normal_batch, lambda x: x, n=0)


def test_monte_carlo_probability_returns_documented_metrics():
    from mcm_toolkit.solve import monte_carlo_probability

    r = monte_carlo_probability(lambda z: z > 0, _normal_batch, n=20000, random_state=0)
    assert r.success
    assert set(r.metrics) == {"probability", "ci_low", "ci_high", "std_error"}
    assert r.data.shape == (20000,)
    assert 0.0 <= r.metrics["ci_low"] <= r.metrics["probability"]
    assert r.metrics["probability"] <= r.metrics["ci_high"] <= 1.0


def test_monte_carlo_probability_std_error_matches_binomial_formula():
    """`std_error` 就是比例的标准误 √(p(1-p)/n)。"""
    from mcm_toolkit.solve import monte_carlo_probability

    n = 20000
    r = monte_carlo_probability(lambda z: z > 0, _normal_batch, n=n, random_state=0)
    p = r.metrics["probability"]
    assert np.isclose(r.metrics["std_error"], np.sqrt(p * (1 - p) / n))


def test_monte_carlo_probability_ci_narrows_with_n():
    """区间宽度随 n 增大而收窄（O(1/√n)）。"""
    from mcm_toolkit.solve import monte_carlo_probability

    def width(n):
        r = monte_carlo_probability(lambda z: z > 0, _normal_batch, n=n, random_state=0)
        return r.metrics["ci_high"] - r.metrics["ci_low"]

    assert width(100_000) < width(1_000)


def test_monte_carlo_probability_higher_confidence_is_wider():
    from mcm_toolkit.solve import monte_carlo_probability

    def width(confidence):
        r = monte_carlo_probability(
            lambda z: z > 0, _normal_batch, n=5000, random_state=0, confidence=confidence
        )
        return r.metrics["ci_high"] - r.metrics["ci_low"]

    assert width(0.99) > width(0.90)


def test_monte_carlo_probability_ci_matches_statsmodels_wilson():
    """区间必须对得上独立实现（statsmodels 的 wilson）—— `^` 写成异或那次就靠它拦住。"""
    from statsmodels.stats.proportion import proportion_confint

    from mcm_toolkit.solve import monte_carlo_probability

    n = 5000
    r = monte_carlo_probability(lambda z: z > 0, _normal_batch, n=n, random_state=0)
    count = int(r.data.sum())
    lo, hi = proportion_confint(count, n, alpha=0.05, method="wilson")
    assert np.isclose(r.metrics["ci_low"], lo)
    assert np.isclose(r.metrics["ci_high"], hi)


def test_monte_carlo_probability_rejects_bad_confidence():
    from mcm_toolkit.solve import monte_carlo_probability

    for bad in (0.0, 1.0, -0.5, 1.5):
        with pytest.raises(ValueError):
            monte_carlo_probability(lambda z: z > 0, _normal_batch, n=100, confidence=bad)


def test_random_walk_shape_and_start():
    """`data` 形状 (n_paths, n_steps + 1)，第 0 列是起点。"""
    from mcm_toolkit.solve import random_walk

    r = random_walk(n_steps=100, n_paths=5, start=2.5, random_state=0)
    assert r.success
    assert r.data.shape == (5, 101)
    assert np.all(r.data[:, 0] == 2.5)


def test_random_walk_default_steps_are_plus_minus_one():
    """默认是对称随机游走：每一步都恰好是 ±1。"""
    from mcm_toolkit.solve import random_walk

    r = random_walk(n_steps=200, n_paths=3, random_state=0)
    assert set(np.unique(np.diff(r.data, axis=1))) == {-1.0, 1.0}


def test_random_walk_is_reproducible():
    from mcm_toolkit.solve import random_walk

    a = random_walk(n_steps=50, n_paths=4, random_state=7)
    b = random_walk(n_steps=50, n_paths=4, random_state=7)
    assert np.array_equal(a.data, b.data)


def test_random_walk_accepts_custom_step_sampler():
    """换成高斯增量就是带漂移的高斯游走 —— 建模差异全在 step_sampler 里。"""
    from mcm_toolkit.solve import random_walk

    r = random_walk(
        n_steps=100,
        n_paths=3,
        step_sampler=lambda rng, size: rng.normal(0.5, 1.0, size),
        random_state=0,
    )
    assert r.success
    assert r.data.shape == (3, 101)
    assert abs(r.data[:, -1].mean() - 50.0) < 20.0   # 100 步 × 漂移 0.5


def test_random_walk_rejects_bad_arguments():
    from mcm_toolkit.solve import random_walk

    with pytest.raises(ValueError):
        random_walk(n_steps=0)
    with pytest.raises(ValueError):
        random_walk(n_steps=10, n_paths=0)
    with pytest.raises(TypeError):
        random_walk(n_steps=10, start="a")  # type: ignore[arg-type]  # 就是要传错


def test_random_walk_docstring_example_runs():
    from mcm_toolkit.solve import random_walk

    r = random_walk(n_steps=100, n_paths=5, random_state=0)
    assert r.data.shape == (5, 101)
    assert np.all(np.abs(np.diff(r.data, axis=1)) == 1)


@pytest.mark.skip(reason=_NOT_IMPLEMENTED)
def test_markov_steady_state_two_state_chain():
    """P=[[0.9,0.1],[0.2,0.8]] 的平稳分布是 [2/3, 1/3]。"""
    from mcm_toolkit.solve import markov_steady_state

    r = markov_steady_state(np.array([[0.9, 0.1], [0.2, 0.8]]))
    assert np.allclose(r.data, [2 / 3, 1 / 3], atol=1e-6)


@pytest.mark.skip(reason=_NOT_IMPLEMENTED)
def test_knapsack_dp_known_optimum():
    from mcm_toolkit.solve import knapsack_dp

    r = knapsack_dp(weights=[1, 3, 4, 5], values=[1, 4, 5, 7], capacity=7)
    assert np.isclose(r.metrics["total_value"], 9.0)
    assert np.isclose(r.metrics["total_weight"], 7.0)


@pytest.mark.skip(reason=_NOT_IMPLEMENTED)
def test_linear_program_basic():
    from mcm_toolkit.solve import linear_program

    r = linear_program(c=[1, 2], A_ub=[[1, 1]], b_ub=[4], bounds=[(0, None)] * 2)
    assert r.success
    assert np.isclose(r.metrics["objective"], 4.0)


@pytest.mark.skip(reason=_NOT_IMPLEMENTED)
def test_mm1_metrics_closed_form():
    from mcm_toolkit.solve import mm1_metrics

    r = mm1_metrics(arrival_rate=2, service_rate=3)
    assert np.isclose(r.metrics["rho"], 2 / 3, rtol=1e-6)
    assert np.isclose(r.metrics["L"], 2.0, rtol=1e-6)


@pytest.mark.skip(reason=_NOT_IMPLEMENTED)
def test_mm1_rejects_unstable_system():
    """λ >= μ 时系统不稳定，应 success=False 而不是返回 inf。"""
    from mcm_toolkit.solve import mm1_metrics

    assert mm1_metrics(arrival_rate=3, service_rate=3).success is False


@pytest.mark.skip(reason=_NOT_IMPLEMENTED)
def test_simulated_annealing_finds_quadratic_minimum():
    from mcm_toolkit.solve import simulated_annealing

    f = lambda v: (v[0] - 3.0) ** 2 + (v[1] + 1.0) ** 2  # noqa: E731
    r = simulated_annealing(f, x0=[0.0, 0.0], n_iter=5000, random_state=0)
    assert np.allclose(r.params["x"], [3.0, -1.0], atol=0.1)