"""蒙特卡洛：估计 π、积 x²、算概率，并画收敛曲线。

三个例子覆盖 MC 的三种典型用法：
积分（几何量）、概率估计、以及"n 要取多大才够"的收敛性验证。
"""

from __future__ import annotations

import numpy as np

from mcm_toolkit.solve import monte_carlo_integral, monte_carlo_probability
from mcm_toolkit.viz import mcm_style, plot_convergence, save_figure


def main() -> None:
    mcm_style()

    # --- 1) 估 π：单位正方形内落入四分之一圆的比例 × 4 ---
    def in_circle(X: np.ndarray) -> np.ndarray:
        return (X**2).sum(axis=1) <= 1.0

    r = monte_carlo_probability(in_circle, lambda rng, **k: rng.uniform(-1, 1, size=(1, 2)), n=200_000, random_state=1)
    print(f"π ≈ {4 * r.metrics['probability']:.5f}  (真值 3.14159)")

    # --- 2) 积分 ∫₀¹ x² dx = 1/3 ---
    r2 = monte_carlo_integral(lambda X: X[:, 0] ** 2, bounds=[(0, 1)], n=200_000, random_state=2)
    print(f"∫₀¹x²dx ≈ {r2.data:.5f}  (真值 0.33333)")

    # --- 3) 收敛性：MC 误差是 O(1/√n)，论文里要给这张图证明 n 选得够 ---
    ns = np.unique(np.round(np.logspace(2, 6, 40)).astype(int))
    errors = []
    for n in ns:
        est = monte_carlo_integral(lambda X: X[:, 0] ** 2, bounds=[(0, 1)], n=int(n), random_state=7)
        errors.append(abs(est.data - 1 / 3))

    fig, _ = plot_convergence(
        np.column_stack([ns, errors]),
        log=True,
        reference=None,
        title="Monte Carlo 误差随样本量收敛",
        xlabel="样本量 n",
        ylabel="绝对误差",
    )
    save_figure(fig, "figures/mc_convergence")


if __name__ == "__main__":
    main()