# /// script
# requires-python = ">=3.12"
# dependencies = ["symbolica==3.0.2", "sympy==1.14.0"]
# ///
"""Expression storage, excluding startup, allocator reserves and workspaces.

uv run scripts/benchmarks/expression-storage.py \
    --output scripts/benchmarks/expression-storage-results.json

Symbolica: packed payload. FORM: sizeof_ words with Off Compress, checked against Bytes used.
Mathematica: live kernel bytes released by deleting the expression after
clearing caches, with/without Share, corrected by an empty-expression control.
SymPy: live traced Python allocation bytes released after clearing caches,
deleting the expression and collecting garbage, corrected by an empty control.
These counters exclude global symbol/type definitions and allocator reserves;
they are expression representation measurements, not total or peak process RAM.
"""

from __future__ import annotations

import argparse
from datetime import datetime, timezone
import gc
import hashlib
from importlib.metadata import version
import json
import math
import os
from pathlib import Path
import platform
import re
import subprocess
import sys
import tempfile
import tracemalloc


DEFAULT_CASES = ["expanded_polynomial", "nested_tree"]
PRODUCT_CASES = ["indexed_functions", "shifted_functions"]
HELD_CASES = ["held_indexed_functions", "held_shifted_functions"]
FACTORED_CASES = ["factored_indexed_functions", "factored_shifted_functions"]
NESTED_CASES = HELD_CASES + FACTORED_CASES
CASES = DEFAULT_CASES + PRODUCT_CASES + NESTED_CASES
ENGINES = ["symbolica", "form", "mathematica", "sympy"]
VARIANTS = {
    "symbols": ["a", "b", "c", "d"],
    "atomic_functions": ["f(a)", "b", "f(c)", "d"],
    "compound_functions": ["f(1+a)", "b", "f(1+c)", "d"],
    "nested_functions": ["f(g(1+a))", "b", "f(g(1+c))", "d"],
    "polynomial_arguments": ["f(1+2*a+2*b+a^2+2*a*b+b^2)", "b",
                             "f(1+2*c+2*d+c^2+2*c*d+d^2)", "d"],
}


def polynomial_source(variant, degree):
    return "(1+" + "+".join(VARIANTS[variant]) + f")^{degree}"


def product_atoms(case, factors):
    return [f"f({'x+' if case.endswith('shifted_functions') else ''}{i})"
            for i in range(1, factors+1)]


def product_factor_sources(case, factors, shared_symbols):
    common = "1" + "".join("+"+name for name in "abc"[:shared_symbols])
    return [f"({common}+{atom})" for atom in product_atoms(case, factors)]


def product_source(case, factors, shared_symbols):
    parts = product_factor_sources(case, factors, shared_symbols)
    return "*".join("g"+p if case in HELD_CASES else p for p in parts)


def split_top_level(source, delimiter):
    parts, depth, start = [], 0, 0
    for index, char in enumerate(source):
        if char == "(":
            depth += 1
        elif char == ")":
            depth -= 1
        elif char == delimiter and depth == 0:
            parts.append(source[start:index])
            start = index+1
    assert depth == 0
    parts.append(source[start:])
    return parts


def verify_held(source, case, factors, shared_symbols, wrapped=True):
    """Check every sum and index, optionally inside FORM's protective g."""
    source = re.sub(r"\s+", "", source).replace("[", "(").replace("]", ")").strip("+;")
    parts = split_top_level(source, "*")
    assert len(parts) == factors
    indices = []
    expected_common = {"1", *"abc"[:shared_symbols]}
    for part in parts:
        prefix = "g(" if wrapped else "("
        assert part.startswith(prefix) and part.endswith(")"), part[:100]
        terms = split_top_level(part[len(prefix):-1], "+")
        calls = [t for t in terms if t.startswith("f(") and t.endswith(")")]
        assert len(calls) == 1 and set(terms)-set(calls) == expected_common
        argument = split_top_level(calls[0][2:-1], "+")
        if case.endswith("shifted_functions"):
            assert len(argument) == 2 and "x" in argument
            argument.remove("x")
        assert len(argument) == 1 and argument[0].isdigit()
        indices.append(int(argument[0]))
    assert sorted(indices) == list(range(1, factors+1))
    return {"terms": 1, "factors": factors, "shared_symbols": shared_symbols,
            "distinct_function_calls": factors, "containers": factors if wrapped else 0,
            "canonical_sha256": digest(product_source(case, factors, shared_symbols))}


def product_terms(factors, shared_symbols):
    return sum(math.comb(factors, k)*math.comb(factors-k+shared_symbols, shared_symbols)
               for k in range(factors+1))


def product_check(case, factors, shared_symbols):
    substitutions = [(atom, i+1) for i, atom in enumerate(product_atoms(case, factors), 1)]
    substitutions += list(zip("abc"[:shared_symbols], [2, 3, 5]))
    common_value = 1 + sum([2, 3, 5][:shared_symbols])
    return substitutions, math.prod(common_value+i+1 for i in range(1, factors+1))


def python_product(engine, case, factors, shared_symbols, namespace):
    if case in NESTED_CASES:
        if engine == "symbolica":
            from symbolica import E
            return E(product_source(case, factors, shared_symbols))
        import sympy as sp
        # Construct the outer product once; each g protects one sum.
        terms = (sp.sympify(factor, locals=namespace)
                 for factor in product_factor_sources(case, factors, shared_symbols))
        return sp.Mul(*(namespace["g"](term) if case in HELD_CASES else term for term in terms))
    if engine == "symbolica":
        from symbolica import E
        expression = E("1")
        for factor in product_factor_sources(case, factors, shared_symbols):
            # Sort after every multiplication, avoiding an unsorted 5**N stream.
            expression = (expression*E(factor)).expand()
    else:
        import sympy as sp
        expression = sp.Integer(1)
        for factor in product_factor_sources(case, factors, shared_symbols):
            expression = sp.expand(expression*sp.sympify(factor, locals=namespace))
    return expression


def wolfram_atom(atom):
    return atom.replace("(", "[").replace(")", "]")


def tree(f, start, leaves):
    if leaves == 1:
        return start
    half = leaves // 2
    return f(tree(f, start, half), tree(f, start+half, half))


def canonical_tree(leaves):
    return tree(lambda a, b: f"f({a},{b})", 1, leaves)


def digest(text):
    return hashlib.sha256(text.encode()).hexdigest()


def benchmark_json(process):
    if process.returncode:
        raise RuntimeError(process.stdout[-4000:] + process.stderr[-2000:])
    for line in process.stdout.splitlines():
        if line.startswith("BENCHMARK:"):
            return json.loads(line.removeprefix("BENCHMARK:"))
    raise RuntimeError(process.stdout[-4000:] + process.stderr[-2000:])


def verify_python(expression, engine, case, leaves, degree, variant, factors=10, shared_symbols=3):
    if case in NESTED_CASES:
        return verify_held(str(expression), case, factors, shared_symbols, wrapped=case in HELD_CASES)
    if case == "nested_tree":
        expected = digest(canonical_tree(leaves))
        actual = digest(str(expression).replace(" ", ""))
        assert actual == expected
        return {"canonical_sha256": actual, "leaves": leaves,
                "function_nodes": leaves-1, "depth": leaves.bit_length()-1}
    if case in PRODUCT_CASES:
        substitutions, expected = product_check(case, factors, shared_symbols)
        if engine == "symbolica":
            from symbolica import E
            value = expression
            for atom, replacement in substitutions:
                value = value.replace(E(atom), replacement)
            value = value.expand()
        else:
            import sympy as sp
            namespace = dict(zip("abcx", sp.symbols("a b c x")))
            namespace["f"] = sp.Function("f")
            value = expression.xreplace({sp.sympify(atom, locals=namespace): replacement
                                        for atom, replacement in substitutions})
        terms = len(expression) if engine == "symbolica" else len(expression.args)
        assert value == expected and terms == product_terms(factors, shared_symbols)
        return {"terms": terms, "factors": factors, "shared_symbols": shared_symbols,
                "distinct_function_calls": factors, "verification_value": expected}
    if engine == "symbolica":
        from symbolica import E
        value = expression
        for atom in VARIANTS[variant]:
            value = value.replace(E(atom), 1)
        value = value.expand()
        terms = len(expression)
    else:
        import sympy as sp
        namespace = dict(zip("abcd", sp.symbols("a b c d")))
        namespace.update(f=sp.Function("f"), g=sp.Function("g"))
        value = expression.xreplace({sp.sympify(atom, locals=namespace): 1
                                    for atom in VARIANTS[variant]})
        terms = len(expression.args)
    assert value == 5**degree and terms == math.comb(degree+4, 4)
    return {"terms": terms, "degree": degree, "value_at_all_generators_one": 5**degree}


def python_worker(engine, case, leaves, degree, variant="compound_functions", factors=10, shared_symbols=3):
    if engine == "symbolica":
        from symbolica import E, S
        f = S("f")
        expression = (python_product(engine, case, factors, shared_symbols, {}) if case in PRODUCT_CASES+NESTED_CASES else
                      tree(f, 1, leaves) if case == "nested_tree" else E(polynomial_source(variant, degree)).expand())
        result = {"storage_bytes": expression.get_byte_size(), "method": "get_byte_size"}
        result.update(verify_python(expression, engine, case, leaves, degree, variant, factors, shared_symbols))
    else:
        import sympy as sp
        from sympy.core.cache import clear_cache
        a, b, c, d, x = sp.symbols("a b c d x")
        f = sp.Function("f")
        namespace = {"a": a, "b": b, "c": c, "d": d, "x": x, "f": f, "g": sp.Function("g")}
        expression = None
        tracemalloc.start()
        # Keep the trace counter above the small-integer cache threshold during
        # the control, so both measurements allocate the same-sized counter.
        counter_padding = bytearray(1024)

        def released_bytes():
            nonlocal expression
            clear_cache()
            gc.collect()
            before = tracemalloc.get_traced_memory()[0]
            expression = None
            gc.collect()
            return before - tracemalloc.get_traced_memory()[0]

        released_bytes()
        control = released_bytes()
        expression = (python_product(engine, case, factors, shared_symbols, namespace) if case in PRODUCT_CASES+NESTED_CASES else
                      tree(f, 1, leaves) if case == "nested_tree"
                      else sp.expand(sp.sympify(polynomial_source(variant, degree), locals=namespace)))
        result = verify_python(expression, engine, case, leaves, degree, variant, factors, shared_symbols)
        released = released_bytes()
        result.update({"storage_bytes": released-control,
                       "released_python_allocation_bytes": released,
                       "empty_control_bytes": control, "method": "tracemalloc release"})
        tracemalloc.stop()
    result.update({"engine": engine, "version": version(engine), "case": case, "verified": True})
    return result


def wolfram_worker(case, leaves, degree, shared, variant="compound_functions", factors=10, shared_symbols=3):
    atoms = [wolfram_atom(atom) for atom in VARIANTS[variant]]
    source = f"tree[1,{leaves}]" if case == "nested_tree" else f"Expand[(1+{'+'.join(atoms)})^{degree}]"
    substitution = ",".join(atom+"->1" for atom in atoms)
    verification = (r'''
        canonical=StringReplace[ToString[expression,InputForm,PageWidth->Infinity],
            {"["->"(","]"->")",WhitespaceCharacter..->""}];
        check=Hash[canonical,"SHA256","HexString"]; Clear[canonical];
    ''' if case == "nested_tree" else f'''
        check=Length[expression]=={math.comb(degree+4,4)} &&
            (expression/.{{{substitution}}})==5^{degree};
        If[!TrueQ[check],Exit[2]];
    ''')
    if case in PRODUCT_CASES:
        # Convert only function parentheses, keeping grouping parentheses intact.
        source = "Expand[" + re.sub(r"f\(([^()]*)\)", r"f[\1]", product_source(case, factors, shared_symbols)) + "]"
        substitutions, expected = product_check(case, factors, shared_symbols)
        substitution = ",".join(f"{wolfram_atom(atom)}->{value}" for atom, value in substitutions)
        verification = f'''
            check=Length[expression]=={product_terms(factors, shared_symbols)} &&
                (expression/.{{{substitution}}})=={expected};
            If[!TrueQ[check],Exit[2]];
        '''
    if case in NESTED_CASES:
        common = "1"+"".join("+"+name for name in "abc"[:shared_symbols])
        argument = "x+i" if case.endswith("shifted_functions") else "i"
        factor = f"{common}+f[{argument}]"
        if case in HELD_CASES:
            factor = f"g[{factor}]"
        source = f"Times@@Table[{factor},{{i,{factors}}}]"
        verification = r'''
            check=StringReplace[ToString[expression,InputForm,PageWidth->Infinity],
                {"["->"(","]"->")",WhitespaceCharacter..->""}];
        '''
    code = r'''
        $HistoryLength=0;
        tree[start_Integer,width_Integer]:=If[width==1,start,
            With[{half=Quotient[width,2]},f[tree[start,half],tree[start+half,half]]]];
        releasedBytes[]:=Module[{before=0,after=0},
            ClearSystemCache[]; before=MemoryInUse[];
            Clear[expression]; after=MemoryInUse[]; before-after];
        expression=0; releasedBytes[]; expression=0; control=releasedBytes[];
        expression=__SOURCE__;
        __VERIFY__
        byteCount=ByteCount[expression];
        shareFreed=__SHARE__;
        released=releasedBytes[];
        Print["BENCHMARK:"<>ExportString[<|
            "engine"->"mathematica", "version"->$Version, "case"->"__CASE__",
            "shared"->__SHARED__, "storage_bytes"->released-control,
            "released_live_kernel_bytes"->released, "empty_control_bytes"->control,
            "share_freed_bytes"->shareFreed, "byte_count_ignoring_sharing"->byteCount,
            "method"->"live kernel release after clearing caches",
            "verification"->check, "verified"->True
        |>,"RawJSON","Compact"->True]];
    '''.replace("__SOURCE__", source).replace("__VERIFY__", verification).replace(
        "__SHARE__", "Share[expression]" if shared else "0").replace(
        "__SHARED__", str(shared)).replace("__CASE__", case)
    result = benchmark_json(subprocess.run(
        ["wolframscript", "-code", code], capture_output=True, text=True, timeout=600))
    if case == "nested_tree":
        assert result["verification"] == digest(canonical_tree(leaves))
    if case in PRODUCT_CASES:
        result.update(factors=factors, shared_symbols=shared_symbols,
                      distinct_function_calls=factors, terms=product_terms(factors, shared_symbols))
    if case in NESTED_CASES:
        result.update(verify_held(result.pop("verification"), case, factors, shared_symbols, wrapped=case in HELD_CASES))
    return result


def form_worker(case, leaves, degree, variant="compound_functions", factors=10, shared_symbols=3):
    with tempfile.TemporaryDirectory(prefix="expression-storage-") as directory:
        work = Path(directory)
        source = canonical_tree(leaves) if case == "nested_tree" else polynomial_source(variant, degree)
        substitution = "\n".join(f"id {atom}=1;" for atom in VARIANTS[variant])
        expected_terms, expected_value = math.comb(degree+4, 4), 5**degree
        multiplications = ""
        if case in PRODUCT_CASES:
            source = "*".join("g"+factor for factor in product_factor_sources(case, factors, shared_symbols))
            multiplications = f"#do STEP=1,{factors}\nid,once,g(z?)=z;\n.sort\n#enddo"
            substitutions, expected_value = product_check(case, factors, shared_symbols)
            substitution = "\n".join(f"id {atom}={value};" for atom, value in substitutions)
            expected_terms = product_terms(factors, shared_symbols)
        if case in NESTED_CASES:
            source = "1"
            held_factors = ["g"+factor for factor in product_factor_sources(case, factors, shared_symbols)]
            # Bound FORM's construction workspace without changing the held expression.
            multiplications = "\n".join("Multiply " + "*".join(held_factors[i:i+32]) + ";\n.sort"
                                       for i in range(0, factors, 32))
            if case in FACTORED_CASES:
                # Merge bounded pieces rather than repeatedly growing one term.
                names = [f"Piece{i}" for i in range(0, factors, 32)]
                statements = [f"Local {name} = " + "*".join(held_factors[i:i+32]) + ";"
                              for name, i in zip(names, range(0, factors, 32))]
                statements.append(".sort")
                level = 0
                while len(names) > 1:
                    statements.append("Drop " + ",".join(names) + ";")
                    merged = []
                    for i in range(0, len(names), 2):
                        name = f"Merge{level}x{i}"
                        statements.append(f"Local {name} = " + "*".join(names[i:i+2]) + ";")
                        merged.append(name)
                    statements.append(".sort")
                    names = merged
                    level += 1
                statements.extend([f"Drop {names[0]};", f"Multiply {names[0]};", ".sort"])
                multiplications = "\n".join(statements)
        (work / "expression.inc").write_text(source + "\n")
        verification = ('#write <tree.txt> "%E",Expression' if case == "nested_tree" else f'''
#$terms=termsin_(Expression);
#write "TERMS:%$",$terms
{substitution}
.sort
#write <value.txt> "%E",Expression
''')
        if case in NESTED_CASES:
            verification = '#write <held.txt> "%E",Expression'
        max_term = max(100000, 64*factors+1024) if case in FACTORED_CASES else 4000000
        number_limit = "#: MaxNumberSize 64" if case in FACTORED_CASES else ""
        buffers = (f"#: ScratchSize {max(16000000, 8*max_term)}\n"
                   f"#: SmallSize {max(32000000, 16*max_term)}\n"
                   f"#: LargeSize {max(64000000, 32*max_term)}\n"
                   f"#: SortIOSize {max(4000000, 4*max_term)}" if case in FACTORED_CASES else "")
        (work / "expression.frm").write_text(f'''
#: MaxTermSize {max_term}
#: WorkSpace 1000000000
#: FunctionLevels 64
{number_limit}
{buffers}
Symbols a,b,c,d,x,z;
CFunction f,g;
Off Compress;
Local Expression =
#include expression.inc
;
.sort
{multiplications}
#$words=sizeof_(Expression);
#write "WORDS:%$",$words
#$held=Expression;
#define MEMORYWORDS "0"
#do BIT=30,0,-1
#if sizeof($held) >= {{`MEMORYWORDS'+2^`BIT'}}
#redefine MEMORYWORDS "{{`MEMORYWORDS'+2^`BIT'}}"
#endif
#enddo
#write "MEMORY_WORDS:%s",`MEMORYWORDS'
#$held=0;
Off Statistics;
Format nospaces;
{verification}
.end
''')
        process = subprocess.run(["form", "-q", "expression.frm"], cwd=work,
                                 capture_output=True, text=True, timeout=300)
        if process.returncode:
            raise RuntimeError(process.stdout[-4000:] + process.stderr[-1000:])
        words = int(re.search(r"WORDS:(\d+)", process.stdout).group(1))
        memory_words = int(re.search(r"MEMORY_WORDS:(\d+)", process.stdout).group(1))
        stats_bytes = int(re.findall(r"Bytes used\s*=\s*(\d+)", process.stdout)[-1])
        assert stats_bytes == words*4 == memory_words*4
        if case == "nested_tree":
            printed = re.sub(r"\s+", "", (work / "tree.txt").read_text()).strip("+;")
            assert digest(printed) == digest(source)
        elif case not in NESTED_CASES:
            terms = int(re.search(r"TERMS:(\d+)", process.stdout).group(1))
            value = re.sub(r"\s+", "", (work / "value.txt").read_text()).strip("+;")
            assert terms == expected_terms and int(value) == expected_value
        return {"engine": "form", "case": case, "storage_bytes": stats_bytes,
                "expression_words": words, "word_bytes": 4,
                "in_memory_dollar_words": memory_words,
                "version": subprocess.run(["form", "-v"], capture_output=True,
                                          text=True, check=True).stdout.strip(),
                "compression_enabled": False,
                "method": "Bytes used with Off Compress, verified against in-memory dollar variable and sizeof_", "verified": True,
                **({"factors": factors, "shared_symbols": shared_symbols,
                    "distinct_function_calls": factors, "terms": expected_terms} if case in PRODUCT_CASES else {}),
                **(verify_held((work / "held.txt").read_text(), case, factors, shared_symbols)
                   if case in NESTED_CASES else {})}


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--worker", choices=["symbolica", "sympy"])
    parser.add_argument("--case", choices=CASES)
    parser.add_argument("--variant", choices=list(VARIANTS), default="compound_functions")
    parser.add_argument("--all-variants", action="store_true",
                        help="Measure all five expanded-expression variations")
    parser.add_argument("--all-products", action="store_true",
                        help="Measure indexed and shifted function products at each factor count")
    parser.add_argument("--all-held-products", action="store_true",
                        help="Keep each factor in g instead of expanding the product")
    parser.add_argument("--all-factored-products", action="store_true",
                        help="Keep sums factored; only FORM uses a protective g")
    parser.add_argument("--factor-counts", nargs="+", type=int, default=[4, 6, 8, 10])
    parser.add_argument("--factors", type=int, default=10)
    parser.add_argument("--shared-symbols", type=int, choices=range(4), default=3)
    parser.add_argument("--engines", nargs="+", choices=ENGINES, default=ENGINES)
    parser.add_argument("--leaves", type=int, default=262144)
    parser.add_argument("--degree", type=int, default=30)
    parser.add_argument("--runs", type=int, default=3)
    parser.add_argument("--output", type=Path, default=Path("expression-storage-results.json"))
    args = parser.parse_args()
    if args.runs < 1 or not 2 <= args.degree <= 35 or not 32 <= args.leaves <= 2**18 or args.leaves & (args.leaves-1):
        parser.error("runs >= 1; degree 2..35; leaves a power of two, 32..262144")
    factor_limit = 100000 if args.all_factored_products or args.case in FACTORED_CASES else 1000 if args.all_held_products or args.case in HELD_CASES else 18
    if not all(2 <= n <= factor_limit for n in [args.factors, *args.factor_counts]):
        parser.error(f"product factor counts must be between 2 and {factor_limit}")
    if sum([args.all_variants, args.all_products, args.all_held_products, args.all_factored_products]) > 1:
        parser.error("choose one batch mode")
    if args.worker:
        print("BENCHMARK:" + json.dumps(python_worker(args.worker, args.case, args.leaves, args.degree,
              args.variant, args.factors, args.shared_symbols)))
        return
    if args.all_products or args.all_held_products or args.all_factored_products:
        product_cases = FACTORED_CASES if args.all_factored_products else HELD_CASES if args.all_held_products else PRODUCT_CASES
        combined = None
        with tempfile.TemporaryDirectory(prefix="indexed-products-") as directory:
            for factors in sorted(set(args.factor_counts)):
                for case in product_cases:
                    output = Path(directory) / f"{case}-{factors}.json"
                    subprocess.run([sys.executable, str(Path(__file__).resolve()), "--case", case,
                                    "--factors", str(factors), "--shared-symbols", str(args.shared_symbols),
                                    "--runs", str(args.runs), "--engines", *args.engines,
                                    "--output", str(output)], check=True)
                    part = json.loads(output.read_text())
                    if combined is None:
                        combined = {k: v for k, v in part.items() if k not in
                                    ["variant", "expressions", "results", "degree", "leaves"]}
                        combined.update({"product_cases": product_cases,
                                         "factor_counts": sorted(set(args.factor_counts)),
                                         "shared_symbols": args.shared_symbols,
                                         "expressions": {c: product_source(c, max(args.factor_counts),
                                                                         args.shared_symbols)
                                                         for c in product_cases}, "results": []})
                    combined["results"].extend(part["results"])
                    args.output.parent.mkdir(parents=True, exist_ok=True)
                    args.output.write_text(json.dumps(combined, indent=2)+"\n")
        return
    if args.all_variants:
        combined = None
        with tempfile.TemporaryDirectory(prefix="expression-variations-") as directory:
            for variant in VARIANTS:
                output = Path(directory) / f"{variant}.json"
                subprocess.run([sys.executable, str(Path(__file__).resolve()),
                                "--case", "expanded_polynomial", "--variant", variant,
                                "--degree", str(args.degree), "--runs", str(args.runs),
                                "--engines", *args.engines, "--output", str(output)], check=True)
                part = json.loads(output.read_text())
                if combined is None:
                    combined = {k: v for k, v in part.items() if k not in ["variant", "expressions", "results"]}
                    combined.update({"variants": VARIANTS, "results": []})
                combined["results"].extend(part["results"])
                args.output.parent.mkdir(parents=True, exist_ok=True)
                args.output.write_text(json.dumps(combined, indent=2)+"\n")
        return
    affinity = min(os.sched_getaffinity(0))
    os.sched_setaffinity(0, {affinity})
    data = {"schema_version": 1, "measured_at_utc": datetime.now(timezone.utc).isoformat(),
            "platform": platform.platform(), "python": platform.python_version(),
            "logical_cpu_affinity": affinity, "runs_per_case": args.runs,
            "metric": "expression representation storage, excluding allocator reserves and workspaces",
            "scope": "Symbolica packed payload; FORM words; Mathematica live kernel allocation release; SymPy traced Python allocation release",
            "variant": args.variant,
            "expressions": {"expanded_polynomial": polynomial_source(args.variant, args.degree),
                            "nested_tree": f"Balanced binary f tree with distinct integer leaves 1..{args.leaves}"},
            "form_encoding": "uncompressed native term representation (Off Compress)",
            "leaves": args.leaves, "degree": args.degree, "results": []}
    if args.case in PRODUCT_CASES+NESTED_CASES:
        data["expressions"] = {args.case: product_source(args.case, args.factors, args.shared_symbols)}
        data.update(shared_symbols=args.shared_symbols)
    for case in [args.case] if args.case else DEFAULT_CASES:
        for run in range(1, args.runs+1):
            for engine in args.engines:
                if engine == "mathematica":
                    results = [wolfram_worker(case, args.leaves, args.degree, share, args.variant,
                                             args.factors, args.shared_symbols)
                               for share in (False, True)]
                elif engine == "form":
                    results = [form_worker(case, args.leaves, args.degree, args.variant,
                                           args.factors, args.shared_symbols)]
                else:
                    result = benchmark_json(subprocess.run(
                        [sys.executable, str(Path(__file__).resolve()), "--worker", engine,
                         "--case", case, "--leaves", str(args.leaves), "--degree", str(args.degree),
                         "--variant", args.variant, "--factors", str(args.factors),
                         "--shared-symbols", str(args.shared_symbols)],
                        capture_output=True, text=True, timeout=900))
                    results = [result]
                for result in results:
                    result["run"] = run
                    if case == "expanded_polynomial":
                        result["variant"] = args.variant
                    data["results"].append(result)
                    suffix = " + Share" if result.get("shared") else ""
                    print(f"{case:20} run {run} {engine+suffix:20} "
                          f"{result['storage_bytes']:,} bytes", flush=True)
                args.output.parent.mkdir(parents=True, exist_ok=True)
                args.output.write_text(json.dumps(data, indent=2)+"\n")


if __name__ == "__main__":
    main()