import json
import sys
import numpy as np
from moment_sharp_bench import u2_from_moments


def math_check(seed=569):
    rng = np.random.default_rng(seed)
    violations, gaps = [], []
    for d in (3, 5, 16, 64):
        for _ in range(100):
            x = np.exp(rng.normal(size=d))
            u = u2_from_moments(x.sum(), (x * x).sum(), d)
            violations.append(max(0.0, float(x.max() - u)))
            gaps.append(float(u - x.max()))
    clustered = np.array([9.0, 2.0, 2.0, 2.0, 2.0])
    exact_error = abs(u2_from_moments(clustered.sum(), (clustered ** 2).sum(), 5) - 9.0)
    return {
        "max_upper_bound_violation": float(max(violations)),
        "mean_bound_gap_squared": float(np.mean(gaps)),
        "clustered_recovery_abs_error": float(exact_error),
        "passed": max(violations) < 1e-9 and exact_error < 1e-9,
    }


def harness_check():
    result = {"available": False, "error": None}
    try:
        sys.path.insert(0, "/home/maxwelhelp/all/math2nn")
        import bench
        result["available"] = True
        result["module"] = getattr(bench, "__file__", None)
        result["symbols"] = [name for name in
                              ("TRACKS", "train_model", "sweep_baseline", "make_report")
                              if hasattr(bench, name)]
    except Exception as exc:
        result["error"] = repr(exc)
    return result


def main():
    report = {
        "stage": "stage-2-bench",
        "math_check": math_check(),
        "harness_check": harness_check(),
        "bench_report": {
            "baseline_sweep": [],
            "idea_results": [],
            "paired_delta": None,
            "permutation_p_value": None,
            "mechanism_signature": None,
            "status": "blocked_missing_bench_checkout",
        },
    }
    with open("stage2_report.json", "w") as f:
        json.dump(report, f, indent=2)
    print(json.dumps(report, indent=2))


if __name__ == "__main__":
    main()