"""Measured scalar/NumPy CPU work; no target execution, providers or network."""

import argparse
import json
import platform
import random
import statistics
import sys
import time
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT / "packages/core"))
from chameleon_core.microfold import compile_model, search, transition_batch, validate_plan


def branching_model(seed):
    rng = random.Random(seed)
    facts = ["f" + str(i) for i in range(17)]
    operations = []
    for i in range(40):
        add = rng.randrange(16)
        remove = (add + rng.randrange(1, 16)) % 16
        operations.append({"id": "operation:" + str(i), "requires": [], "forbids": [],
                           "adds": [facts[add]], "removes": [facts[remove]], "cost": 1 + i % 3})
    return {"version": "microfold-model-v1", "evidence_digest": "b" * 64, "facts": facts,
            "initial": [], "goal": {"all": [facts[16]], "none": []}, "operators": operations}


def measure(work, repeats):
    timings = []
    for _ in range(repeats):
        start = time.perf_counter_ns()
        work()
        timings.append((time.perf_counter_ns() - start) / 1_000_000)
    ordered = sorted(timings)
    return {"median_ms": statistics.median(timings), "p95_ms": ordered[int((len(ordered)-1)*.95)],
            "min_ms": min(timings), "repeats": repeats}


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--output", required=True)
    parser.add_argument("--repeats", type=int, default=31)
    args = parser.parse_args()
    if not 3 <= args.repeats <= 101:
        raise ValueError("BENCHMARK_REPEAT_BUDGET")
    destination = Path(args.output).resolve()
    if destination.exists() or not destination.is_relative_to(ROOT / ".local"):
        raise ValueError("FRESH_LOCAL_REPORT_REQUIRED")
    started = time.perf_counter_ns()
    import numpy as np

    cold_ms = (time.perf_counter_ns() - started) / 1_000_000
    model = branching_model(20260921)
    _, _, _, operations = compile_model(model)
    rng = random.Random(27)
    batches = [[rng.randrange(1 << 16) for _ in range(25)] for _ in range(10)]
    # Ten different 25-state x 40-operator batches: exactly 10,000 pair checks.
    for states in batches:
        assert transition_batch(states, operations, "python") == transition_batch(states, operations, "numpy")
    parity = 0
    for seed in range(25):
        source = branching_model(seed)
        if seed % 2:
            source["goal"]["all"] = ["f1", "f2", "f3"]
        scalar, vector = search(source, backend="python"), search(source, backend="numpy")
        assert {k: v for k, v in scalar.items() if k not in {"backend", "wall_ns"}} == {
            k: v for k, v in vector.items() if k not in {"backend", "wall_ns"}}
        if scalar["status"] == "PLAN_FOUND":
            assert validate_plan(source, scalar) and validate_plan(source, vector)
        parity += 1
    reports = {}
    for backend in ("python", "numpy"):
        def kernel(backend=backend):
            return [transition_batch(states, operations, backend) for states in batches]

        kernel()  # Warm runtime separately from measured iterations.
        full = search(model, backend=backend)
        reports[backend] = {"kernel_10000": measure(kernel, args.repeats),
                            "full_search": measure(lambda backend=backend: search(model, backend=backend), args.repeats),
                            "full_search_example": full}
    result = {"version": "microfold-benchmark-v1", "python": platform.python_version(),
              "platform": platform.platform(), "numpy": np.__version__, "numpy_import_ms": cold_ms,
              "parity_models": parity, "kernel_microsteps_per_sample": 10000, "backends": reports,
              "external_steps": 0,
              "scope": "Synthetic explicit transition search; timings exclude browser, acquisition, model inference and network"}
    destination.write_text(json.dumps(result, indent=2), encoding="utf-8")
    print(json.dumps({"report": str(destination), "numpy_import_ms": cold_ms,
                      "parity_models": parity, "backends": {
                          backend: {key: value for key, value in row.items() if key != "full_search_example"}
                          for backend, row in reports.items()}}, indent=2))


if __name__ == "__main__":
    main()