"""Build and verify the F13 or nested analytic expression from the benchmark page.

Install the benchmark's Symbolica source wheel and NumPy, then place this script,
native-output-reference.json, f13.txt and/or large_nested.txt in one directory.
"""
import argparse
import hashlib
import json
from pathlib import Path
import time

import numpy as np
from symbolica import E, S, T


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--case", choices=["f13", "nested", "nested-expanded"], default="nested")
    args = parser.parse_args()
    root = Path(__file__).resolve().parent
    references = json.loads((root / "native-output-reference.json").read_text())
    reference = references["f13-expanded" if args.case == "f13" else "nested-original"]
    source = (root / ("f13.txt" if args.case == "f13" else "large_nested.txt")).read_bytes()
    assert hashlib.sha256(source).hexdigest() == reference["input_sha256"]
    started = time.perf_counter()
    expression = E(source.decode())
    if args.case == "nested-expanded":
        argument = S("paired_expand_argument_")
        for function in S("exp", "sin", "cos"):
            expression = expression.replace(function(argument),
                function(argument.hold(T().expand(via_poly=True))), bottom_up=True, nested=True)
        expression = expression.expand(via_poly=True)
    parameters = [S(name) for name in reference["names"]]
    evaluator = expression.evaluator(parameters, iterations=10, n_cores=1,
                                    horner_use_cpe=True, cpe_iterations=None, jit_compile=False)
    evaluator.jit_compile(True, direct_translation=True, optimization_level=3)
    points = np.ascontiguousarray(np.array(reference["input_numerators"], dtype=np.float64)
                                 / reference["input_denominator"])
    actual = evaluator.evaluate(points).ravel()
    elapsed = time.perf_counter() - started
    assert len(actual) == len(reference["expected"]) == 32
    for value, expected in zip(actual, reference["expected"]):
        error = abs(value-expected) / abs(expected) if expected else abs(value)
        assert np.isfinite(value) and error < 1e-9, (value, expected)
    print(f"Built and verified at {len(actual)} points in {elapsed:.3f} s.")
    print("First value:", actual[0])


if __name__ == "__main__":
    main()