"""Benchmark the regional optimizer against verified native FORM programs.

Use with compare-expanded-native.py and count-native-output.py beside this file.
All arguments accepted by compare-expanded-native.py remain available. Additional
options select the original or expanded nested input and Horner CPE scoring.
The helper verifies the input and native binaries against their result manifests.
"""
import argparse
import importlib.util
from pathlib import Path
import sys
import time

HERE = Path(__file__).resolve().parent
HELPER = HERE / "compare-expanded-native.py"
spec = importlib.util.spec_from_file_location("native_output_helpers", HELPER)
helper = importlib.util.module_from_spec(spec)
spec.loader.exec_module(helper)


def build_candidate(args, source, reference, points, provenance):
    from symbolica import E, S, Evaluator

    assert args.policy == "automatic" and not args.horner_order and not args.horner_order_json
    def profile_gate(phase):
        if extra.profile_build:
            (args.output / f"{phase}.ready").touch()
            while not (args.output / f"{phase}.start").exists():
                time.sleep(.002)

    profile_gate("build")
    started = time.perf_counter()
    expression = E(source.read_text())
    parsed = time.perf_counter()
    if args.dataset == "nested" and extra.input_form == "expanded":
        expression = helper.expand_all(expression)
    expanded = time.perf_counter()
    params = [S(name) for name in reference["names"]]
    profile_gate("generation")
    ordered = time.perf_counter()
    settings = dict(iterations=args.iterations, n_cores=1, cpe_iterations=None,
                    max_common_pair_cache_entries=args.pair_cache, jit_compile=False,
                    direct_translation=True, horner_use_cpe=extra.horner_use_cpe == "true")
    evaluator = expression.evaluator(params, **settings)
    prepared = time.perf_counter()
    jit_settings = dict(direct_translation=True, optimization_level=3)
    evaluator.jit_compile(True, **jit_settings)
    evaluator.evaluate(points[:4])
    finished = time.perf_counter()
    if extra.profile_build:
        (args.output / "build.done").touch()
        while not (args.output / "build.ack").exists():
            time.sleep(.002)
    counts = helper.count_instructions(evaluator, len(params))
    verification = helper.check(evaluator.evaluate(points).ravel(), reference["expected"])
    saved = args.output / "evaluator.bin"
    saved.write_bytes(bytes(evaluator.save()))
    restored = Evaluator.load(saved.read_bytes())
    reloaded = helper.check(restored.evaluate(points).ravel(), reference["expected"])
    record = dict(
        label=args.label, provenance=provenance, settings=settings, jit_settings=jit_settings,
        input_form=extra.input_form if args.dataset == "nested" else "expanded polynomial",
        ordering={"policy": "automatic"}, parse_seconds=parsed-started,
        expansion_seconds=expanded-parsed, order_selection_seconds=ordered-expanded,
        optimization_seconds=prepared-ordered, native_compilation_seconds=finished-prepared,
        total_build_seconds=finished-started, build_timeout_seconds=None,
        saved_evaluator=str(saved), saved_evaluator_sha256=helper.digest(saved),
        verification=verification, serialization_verification=reloaded,
        driver_source_sha256=helper.digest(__file__), helper_source_sha256=helper.digest(HELPER),
        **counts,
    )
    return restored, record


if __name__ == "__main__":
    parser = argparse.ArgumentParser(add_help=False)
    parser.add_argument("--input-form", choices=["original", "expanded"], default="original")
    parser.add_argument("--horner-use-cpe", choices=["false", "true"], default="true")
    parser.add_argument("--profile-build", action="store_true",
                        help="Wait for the external memory sampler at build and generation boundaries")
    extra, remaining = parser.parse_known_args()
    sys.argv = [sys.argv[0], *remaining]
    helper.build_candidate = build_candidate
    helper.main()