MCM-ICM_Toolkit/src/mcm_toolkit/solve/simulate.py

"""随机模拟：Monte Carlo / 马尔可夫 / 随机游走 / 排队。

★ 本模块的 `monte_carlo` 与 `monte_carlo_integral` 是 S 档，建议自己实现。
原建议里的 `brownian_motion` `geometric_brownian_motion` 已删（niche），
`simulate_queue` 也已删（等于自己写一个离散事件仿真引擎，成本远高于收益）。
"""

from __future__ import annotations

from collections.abc import Callable
from typing import Any

import numpy as np
from scipy.stats import norm

from mcm_toolkit.core import Result, get_rng

__all__ = [
    "markov_simulate",
    "markov_steady_state",
    "mm1_metrics",
    "monte_carlo",
    "monte_carlo_integral",
    "monte_carlo_probability",
    "random_walk",
]


def monte_carlo(
    sampler: Callable[..., Any],
    simulator: Callable[..., Any],
    n: int | np.integer = 10000,
    random_state: int | None = None,
    **kwargs: Any,
) -> Any:
    """通用蒙特卡洛：采样 → 仿真 → 汇总。

    Parameters
    ----------
    sampler:
        `f(rng, size, **kwargs) -> 一批参数`。**一次就要抽够 `n` 组**，
        例如 `lambda rng, size: rng.normal(0, 1, size)`。
        批量大小由本函数以 `size=n` 传进去：`sampler` 必须能接住 `size`
        （显式声明，或带 `**kwargs`）。**别再往 `**kwargs` 里塞 `size`**，否则重复传参。
    simulator:
        `f(params) -> 一批结果`。收到 `sampler` 的那一整批参数，
        返回长度同为 `n` 的数组（标量 / 长度对不上都算契约错误）。
        **内部不要写 Python 循环**。
    n:
        模拟次数，正整数。与 `simulator` 返回的长度必须一致。
    random_state:
        随机种子。`None` 时跟随 `core.set_seed` 的全局流，见 `core.random.get_rng`。
    **kwargs:
        原样转交给 `sampler`，例如分布参数 `mu` / `sigma`（`size` 由本函数自己传，见上）。
        会一并记进 `Result.params`，便于复现。

    Returns
    -------
    `Result`：`data` 是 `n` 个结果组成的数组（一律转成 float，布尔结果按 0/1 统计），
    `metrics={"mean", "std", "median", "p05", "p95", "ci_low", "ci_high"}`。

    Raises
    ------
    TypeError:
        `sampler` / `simulator` 不可调用，`n` 不是整数，或 `random_state` 类型非法。
    ValueError:
        `n` 不是正整数。

    Notes
    -----
    比赛里 Monte Carlo 的收敛速度是 `O(1/√n)` —— 想少一半误差就要 4 倍样本。
    所以论文里要给"区间随 n 收敛"的图，证明 `n` 已经够大
    （这也是 `viz.plot_convergence` 的典型用法）。

    结果里只要出现 **NaN** 就判失败（`success=False`）：MC 的 NaN 通常来自除零、
    取负数的根号或模型发散，静默统计只会得到一列 nan 指标。

    Examples
    --------
    >>> import numpy as np
    >>> draw = lambda rng, size, **k: rng.normal(0, 1, size)
    >>> r = monte_carlo(draw, lambda x: x ** 2, n=10000, random_state=0)
    >>> bool(abs(r.metrics["mean"] - 1.0) < 0.1)  # E[X²] = 1
    True
    """
    if not callable(sampler) or not callable(simulator):
        raise TypeError("sampler 和 simulator 必须是可调用函数")
    if isinstance(n, bool) or not isinstance(n, (int, np.integer)):
        raise TypeError(f"n 必须是整数，你传入了 {type(n).__name__}")
    if int(n) <= 0:
        raise ValueError(f"n 必须是正整数，你传入了 {n}")
    n = int(n)

    rng = get_rng(random_state)
    call_params = {"n": n, "random_state": random_state, **kwargs}

    try:
        params = sampler(rng, size=n, **kwargs)
        data = np.asarray(simulator(params), dtype=float).ravel()
        if np.isnan(data).any():
            raise ValueError("模拟结果中包含 NaN，请检查 sampler 或 simulator 的传入参数")
        if len(data) != n:
            raise ValueError(f"simulator 返回的结果数量({len(data)})必须与模拟次数 n({n})一致")

        metrics = {
            "mean": float(np.mean(data)),
            "std": float(np.std(data, ddof=0)),
            "median": float(np.median(data)),
            "p05": float(np.percentile(data, 5)),
            "p95": float(np.percentile(data, 95)),
            "ci_low": float(np.percentile(data, 2.5)),
            "ci_high": float(np.percentile(data, 97.5)),
        }

        return Result(
            name="monte_carlo",
            params=call_params,
            metrics=metrics,
            data=data,
            success=True,
            message=f"蒙特卡洛模拟成功：已完成 {n} 次模拟",
        )
    except Exception as e:
        return Result(
            name="monte_carlo",
            params=call_params,
            metrics={},
            data=None,
            success=False,
            message=f"蒙特卡洛模拟失败：{type(e).__name__} - {e}",
        )


def monte_carlo_probability(
    event: Callable[..., Any],
    sampler: Callable[..., Any],
    n: int | np.integer = 10000,
    confidence: float = 0.95,
    random_state: int | None = None,
) -> Any:
    """估计事件发生概率，并给出区间。

    Parameters
    ----------
    event:
        `f(params) -> 一批 bool`。吃到 `sampler` 的**整批**参数，
        返回长度为 `n` 的布尔数组（写成 `z > 0` 这种逐元素判断即可，不要写 Python 循环）。
    sampler:
        `f(rng, size, **kwargs) -> 一批参数`，契约同 `monte_carlo`：
        批量大小由本函数以 `size=n` 传入，`sampler` 必须能接住 `size`。
    n:
        模拟次数，正整数。与 `event` 返回的长度必须一致。
    confidence:
        区间置信水平，必须落在开区间 `(0, 1)` 内。
    random_state:
        随机种子。`None` 时跟随 `core.set_seed` 的全局流，见 `core.random.get_rng`。

    Returns
    -------
    `Result`：`data` 是 `n` 个 0/1 结果组成的数组，
    `metrics={"probability", "ci_low", "ci_high", "std_error"}`。
    区间是 Wilson score 区间（`p` 贴近 0 / 1 时比 Wald 稳），并夹在 `[0, 1]` 内（比例量纲，见决策 42）。

    Raises
    ------
    TypeError:
        `n` 不是整数，或 `random_state` 类型非法。
    ValueError:
        `n` 不是正整数，或 `confidence` 不在 `(0, 1)` 内。

    Notes
    -----
    概率估计的相对误差 ≈ `1/√(n·p)`。**小概率事件（p < 1e-3）用朴素 MC 是不可行的**，
    需要重要性采样或方差缩减——遇到这类题要提前意识到，别硬跑。
    本函数当前的口径：`probability <= 1e-3` 时**直接判 `success=False`**（只给提示，不给区间）。

    Examples
    --------
    >>> import numpy as np
    >>> draw = lambda rng, size, **k: rng.normal(0, 1, size)
    >>> r = monte_carlo_probability(lambda z: z > 0, draw, n=2000, random_state=0)
    >>> abs(r.metrics["probability"] - 0.5) < 0.1
    True
    """
    if not callable(sampler) or not callable(event):
        raise TypeError("sampler 和 event 必须是可调用函数")
    if isinstance(n, bool) or not isinstance(n, (int, np.integer)):
        raise TypeError(f"n 必须是整数，你传入了 {type(n).__name__}")
    if int(n) <= 0:
        raise ValueError(f"n 必须是正整数，你传入了 {n}")
    n = int(n)
    if confidence <= 0.0 or confidence >= 1:
        raise ValueError(f"confidence必须在0和1之间,而你输入了{confidence}")

    rng = get_rng(random_state)
    call_params = {"n": n, "random_state": random_state, "confidence": confidence}

    try:
        params = sampler(rng, size=n)
        data = np.asarray(event(params), dtype=float).ravel()
        if np.isnan(data).any():
            raise ValueError("模拟结果中包含 NaN，请检查 event 或 sampler 的传入参数")
        if len(data) != n:
            raise ValueError(f"event 返回的结果数量({len(data)})必须与模拟次数 n({n})一致")
        p_value = float(np.mean(data))
        std_error = float(np.sqrt(p_value * (1.0 - p_value) / n))

        z_score = norm.ppf(1 - (1 - confidence) / 2)
        center = (p_value + z_score**2 / (2 * n)) / (1 + z_score ** 2 / n)
        half_width = (z_score / (1 + z_score ** 2 / n)) * np.sqrt(p_value * (1 - p_value) / n + z_score ** 2 / (4 * n ** 2))
        ci_low = max(0.0, center - half_width)
        ci_high = min(1.0, center + half_width)

        warning_message = ""
        if p_value <= 1e-3 or p_value >= 1 - 1e-3:
            warning_message = "小概率事件 P < 1e-3 或者高概率事件 P > 1 - 1e-3,不建议使用monte_carlo_probability"

        metrics = {
            "probability": p_value,
            "ci_low": float(ci_low),
            "ci_high": float(ci_high),
            "std_error": std_error,
        }
        return Result(
            name="monte_carlo_probability",
            params=call_params,
            metrics=metrics,
            data=data,
            success=True,
            message=(
                f"蒙特卡洛模拟成功：已完成 {n} 次模拟"
                + (f"；{warning_message}" if warning_message else "")
            ),
        )
    except Exception as e:
        return Result(
            name="monte_carlo_probability",
            params=call_params,
            metrics={},
            data=None,
            success=False,
            message=f"蒙特卡洛模拟失败：{type(e).__name__} - {e}",
        )


def monte_carlo_integral(
    func: Callable[..., Any],
    bounds: Any,
    n: int | np.integer = 10000,
    random_state: int | None = None,
) -> Any:
    """用蒙特卡洛求多重积分。

    Parameters
    ----------
    func:
        `f(X) -> 一批函数值`。`X` 是形状 `(n, d)` 的样本矩阵（一行一个样本），
        返回长度 `n` 的数组（写成 `X[:, 0] ** 2` 这种逐行写法，不要写 Python 循环）。
    bounds:
        `[(lo, hi), ...]`（非空 list / tuple），维度即积分重数；要求每项都是二元组且 `lo < hi`。
    n:
        模拟次数，正整数。
    random_state:
        随机种子。`None` 时跟随 `core.set_seed` 的全局流，见 `core.random.get_rng`。

    Returns
    -------
    `Result`：**`metrics["mean"]` 才是积分估计值** ≈ `∫f dV`；
    `data` 是 `n` 个样本各自的贡献（`func(X_i) × 体积`），可拿去画随 `n` 的收敛图。
    `params` 里额外记了 `bounds` 与 `dimension`。

    Raises
    ------
    ValueError:
        `bounds` 不是非空 list / tuple、形状不是 `(d, 2)`（例如漏括号写成 `[0, 1]`），或某个 `lo >= hi`。

    Notes
    -----
    高维积分是 MC 相对网格法的**唯一绝对优势**：网格法误差 `O(h^{-d})` 随维度爆炸，
    MC 始终是 `O(1/√n)`。三维以上优先考虑 MC。

    Examples
    --------
    >>> import numpy as np
    >>> r = monte_carlo_integral(lambda X: X[:, 0] ** 2, bounds=[(0, 1)], n=200000, random_state=0)
    >>> abs(r.metrics["mean"] - 1 / 3) < 0.01   # ∫₀¹ x² dx = 1/3
    True
    """
    if not isinstance(bounds, (list, tuple)) or len(bounds) == 0:
        raise ValueError("bounds 必须是非空的列表或元组")

    bounds_arr = np.asarray(bounds)

    if bounds_arr.ndim != 2 or bounds_arr.shape[1] != 2:
        raise ValueError(f"bounds 必须形如 [(lo, hi), ...]，当前形状为 {bounds_arr.shape}，请检查是否漏了括号")

    lows = bounds_arr[:, 0]
    highs = bounds_arr[:, 1]
    d = len(lows)
    if np.any(lows >= highs):
        raise ValueError("积分边界非法：下界必须严格小于上界")

    widths = highs - lows

    volume = np.prod(widths)

    def sampler(rng, size=n):
        return rng.uniform(low=lows, high=highs, size=(size, d))

    def simulator(X):
        vals = func(X)
        if len(vals) != n:
            raise ValueError(f"func 返回的结果数量({len(vals)})必须与模拟次数 n({n})一致")
        return np.asarray(vals, dtype=float).ravel() * volume

    result = monte_carlo(
        sampler=sampler,
        simulator=simulator,
        n=n,
        random_state=random_state,
    )

    if result.success:
        result.name = "monte_carlo_integral"
        result.params["bounds"] = bounds
        result.params["dimension"] = d
        result.message = f"蒙特卡洛多重积分成功：维度 {d}，体积 {volume:.4f}，模拟 {n} 次。"
    else:
        result.name = "monte_carlo_integral"
        result.params["bounds"] = bounds
        result.params["dimension"] = d
        result.message = f"蒙特卡洛多重积分失败：维度 {d}，体积 {volume:.4f}，模拟 {n} 次。{result.message}"

    return result


def random_walk(
    n_steps: int,
    n_paths: int = 1,
    step_sampler: Callable[..., Any] | None = None,
    start: float = 0.0,
    random_state: int | None = None,
) -> Any:
    """随机游走轨迹。

    Parameters
    ----------
    step_sampler:
        `f(rng, size) -> 增量`，默认为 ±1 等概率（对称随机游走）。

    Returns
    -------
    `Result`：`data` 形状 `(n_paths, n_steps + 1)`。

    Examples
    --------
    >>> r = random_walk(n_steps=100, n_paths=5, random_state=0)
    >>> r.data.shape
    (5, 101)
    """
    # TODO(L1)
    raise NotImplementedError


def markov_simulate(
    transition: Any,
    initial: Any,
    steps: int = 100,
    n_paths: int = 1,
    random_state: int | None = None,
) -> Any:
    """马尔可夫链模拟，返回状态序列。"""
    # TODO(L1)
    raise NotImplementedError


def markov_steady_state(transition: Any, max_iter: int = 10000, tol: float = 1e-10) -> Any:
    """平稳分布。

    Parameters
    ----------
    transition:
        行随机矩阵（每行和为 1）。

    Returns
    -------
    `Result`：`data` 是平稳分布向量，`metrics={"converged", "iterations"}`。

    Notes
    -----
    用迭代法而不是特征分解：迭代法能顺便告诉你"收敛快不快"，
    而收敛速度本身就是论文里值得写一句的性质。
    **不可约 + 非周期**才有唯一平稳分布，周期链会不收敛，要检查。

    Examples
    --------
    >>> import numpy as np
    >>> P = np.array([[0.9, 0.1], [0.2, 0.8]])
    >>> r = markov_steady_state(P)
    >>> bool(np.allclose(r.data, [2 / 3, 1 / 3], atol=1e-6))
    True
    """
    # TODO(L1)
    raise NotImplementedError


def mm1_metrics(arrival_rate: float, service_rate: float) -> Any:
    """M/M/1 排队系统的闭式指标。

    Parameters
    ----------
    arrival_rate, service_rate:
        到达率 λ 与服务率 μ。`λ >= μ` 时系统不稳定，队列长度无界——
        这种情况下函数返回 `success=False` 并说明原因，而不是返回 `inf`。

    Returns
    -------
    `Result`：`metrics={"rho", "L", "Lq", "W", "Wq"}`
    （系统利用率、系统内平均顾客数、队列长度、平均逗留时间、平均等待时间）。

    Notes
    -----
    M/M/1 是唯一值得手写的排队模型：公式是闭式的，10 行搞定。
    多服务员（M/M/c）和一般服务时间分布（M/G/1）在美赛出现频率极低，
    真遇到时现场用仿真做，不做通用封装。

    Examples
    --------
    >>> r = mm1_metrics(arrival_rate=2, service_rate=3)
    >>> round(r.metrics["rho"], 4)
    0.6667
    """
    # TODO(L1)
    raise NotImplementedError