"""Matched Python-wheel benchmark used by wasm_vs_native.mjs.

The benchmark intentionally measures operations on already-constructed
expressions except for ``parse``. Each case has at least five seconds of warmup
followed by nine calibrated batches, targeting 120 ms per batch. Raw batch
averages and their median are reported in nanoseconds per operation.
"""

import gc
import hashlib
import json
import math
import os
import sys
import statistics
import time

import symbolica
from symbolica import E, S


def measure(function, warmup_seconds=5.0, sample_seconds=0.12, samples=9, minimum_operations=0):
    def batch(count):
        start = time.perf_counter_ns()
        for _ in range(count):
            function()
        return time.perf_counter_ns() - start

    count = 1
    while batch(count) < 5_000_000:
        count *= 2
    start = time.perf_counter()
    warmup_ops = 0
    while time.perf_counter() - start < warmup_seconds or warmup_ops < minimum_operations:
        elapsed = batch(count)
        warmup_ops += count
    warmup_elapsed = time.perf_counter() - start
    count = max(1, min(1_000_000, int(sample_seconds * 1e9 * count / elapsed)))
    timings = []
    for _ in range(samples):
        timings.append(batch(count) / count)
    return {"median_ns": statistics.median(timings), "samples_ns": timings,
            "iterations_per_sample": count, "warmup_seconds": warmup_elapsed,
            "warmup_operations": warmup_ops}


def main(case_names=None, warmup_seconds=5.0, minimum_operations=0):
    x, y, z = S("x", "y", "z")
    expression = E("(x+y+1)^30")
    large = E("(1+x+y+z)^30")
    expanded = expression.expand()
    large_expanded = large.expand()
    text = expanded.format_plain()
    assert E(text) == expanded
    factor_input = E("x^12-y^12")
    derivative_input = E("(x^2+y^2)^20")
    evaluation_input = E("(x+y)^12")
    point = {x: 0.1, y: 0.2}
    # Exact inputs from the multiplication/division cases in the FLINT suite.
    left = E("(1+x+y+z)^12").expand().to_polynomial()
    right = E("(1+2*x-y+3*z)^11").expand().to_polynomial()
    divisor = E("(1+2*x-y+3*z)^7").expand().to_polynomial()
    dividend = left * divisor
    assert dividend // divisor == left
    assert expanded.to_polynomial().nterms() == math.comb(32, 2)
    assert large_expanded.to_polynomial().nterms() == math.comb(33, 3)
    assert factor_input.factor().expand() == factor_input
    assert derivative_input.derivative(x) == E("40*x*(x^2+y^2)^19")
    assert abs(evaluation_input.evaluate({x: 2.0, y: 3.0}) - 244140625) < 1e-6
    assert abs(expanded.evaluate(point) - 1.3**30) / 1.3**30 < 1e-12
    cases = {
        "parse_small": lambda: E("(x+y+1)^30"),
        "expand_small": expression.expand,
        "factor_small": factor_input.factor,
        "derivative_small": lambda: derivative_input.derivative(x),
        "evaluate_small": lambda: evaluation_input.evaluate({x: 2.0, y: 3.0}),
        "parse_496_terms": lambda: E(text),
        "expand_5456_terms": large.expand,
        "derivative_5456_terms": lambda: large_expanded.derivative(x),
        "evaluate_496_terms": lambda: expanded.evaluate(point),
        "flint_dense_multiply": lambda: left * right,
        "flint_dense_exact_divide": lambda: dividend // divisor,
    }
    gc.collect()
    results = {}
    for name, function in cases.items():
        if case_names is not None and name not in case_names:
            continue
        print("BENCHMARK " + name, file=sys.stderr, flush=True)
        # This case showed a late V8 tier transition in the two-second pilot.
        results[name] = measure(function, warmup_seconds=max(warmup_seconds, 10.0) if name == "expand_small" else warmup_seconds,
                                minimum_operations=minimum_operations)
    print(json.dumps({"version": symbolica.get_version(), "python": sys.version,
                      "module": symbolica.__file__, "validation": "passed",
                      "input_sha256": hashlib.sha256(text.encode()).hexdigest(),
                      "results": results}, sort_keys=True), flush=True)


if __name__ == "__main__":
    main(**json.loads(os.environ.get("SYMBOLICA_BENCH_OPTIONS", "{}")))