# /// script
# requires-python = ">=3.12"
# dependencies = ["symbolica==3.0.2", "numpy==2.4.1", "mpmath==1.3.0"]
# ///
"""Native evaluation of large_nested.txt, retaining nested elementary functions.

Place f13-output.py beside this script for the shared process profiler and
Wolfram compiler-driver adapter. Run --smoke first, then run the full input.
All timings include expression parsing and any required expansion, but exclude
engine startup, imports, reference evaluation and instruction inspection.
"""
from __future__ import annotations

import argparse
import ast
from collections import Counter
import ctypes
from datetime import datetime, timezone
import hashlib
import importlib.util
from importlib.metadata import version
import json
import math
import os
from pathlib import Path
import platform
import random
import re
import signal
import statistics
import subprocess
import sys
import time

HELPER = Path(__file__).with_name('f13-output.py')
spec = importlib.util.spec_from_file_location('native_benchmark_common', HELPER)
common = importlib.util.module_from_spec(spec)
spec.loader.exec_module(common)
NAMES = [f'x{i}' for i in range(7)]
DENOMINATOR = 16
ROUTES = ['symbolica_asm', 'symbolica_symjit', 'mathematica_unoptimized',
          'mathematica_optimized', 'form_cpp', 'form_cpp_o2']
SMOKE = '(exp(-(x0+x1)^2/4)*(1+x0+x1*x2)^2+sin(x3+x4)^2*cos(x5+x6))'
TOKEN = re.compile(r'\s*(x[0-6]|exp|sin|cos|\d+|[()+*/^\-])')


def rpn(source):
    """Independent, iterative parser; never expands the input expression."""
    tokens = []; offset = 0
    while offset < len(source):
        match = TOKEN.match(source, offset)
        if not match:
            if not source[offset:].strip(): break
            raise ValueError(f'Unsupported syntax at {offset}')
        tokens.append(match[1]); offset = match.end()
    output = []; stack = []; previous = None
    precedence = {'+': 1, '-': 1, '*': 2, '/': 2, 'neg': 3, 'pos': 3, '^': 4}
    for token in tokens:
        if token.isdigit() or token in NAMES: output.append(token)
        elif token in {'exp', 'sin', 'cos', '('}: stack.append(token)
        elif token == ')':
            while stack and stack[-1] != '(': output.append(stack.pop())
            assert stack and stack.pop() == '('
            if stack and stack[-1] in {'exp', 'sin', 'cos'}: output.append(stack.pop())
        else:
            unary = token in {'+', '-'} and (previous is None or previous in precedence or previous == '(')
            op = ('neg' if token == '-' else 'pos') if unary else token
            if not unary:
                while stack and stack[-1] in precedence and (
                    precedence[stack[-1]] > precedence[op] or
                    (precedence[stack[-1]] == precedence[op] and op != '^')):
                    output.append(stack.pop())
            stack.append(op)
        previous = token
    while stack:
        token = stack.pop(); assert token != '('; output.append(token)
    return output


def evaluate_reference(program, point, precision):
    import mpmath as mp
    with mp.workdps(precision):
        values = dict(zip(NAMES, [mp.mpf(n)/DENOMINATOR for n in point]))
        stack = []
        for token in program:
            if token in values: stack.append(values[token])
            elif token.isdigit(): stack.append(mp.mpf(token))
            elif token in {'exp', 'sin', 'cos'}: stack[-1] = getattr(mp, token)(stack[-1])
            elif token == 'neg': stack[-1] = -stack[-1]
            elif token == 'pos': pass
            else:
                b = stack.pop(); a = stack.pop()
                stack.append({'+': lambda: a+b, '-': lambda: a-b, '*': lambda: a*b,
                              '/': lambda: a/b, '^': lambda: a**b}[token]())
        assert len(stack) == 1 and mp.isfinite(stack[0])
        return mp.nstr(stack[0], precision)


def reference_data(source, batch_size):
    import mpmath as mp
    program = rpn(source)
    rng = random.Random(1729)
    points = [[rng.randint(-12, 12) for _ in NAMES] for _ in range(32)]
    expected = []; exact = []
    for point in points:
        a = evaluate_reference(program, point, 80)
        b = evaluate_reference(program, point, 120)
        with mp.workdps(120): assert abs(mp.mpf(a)-mp.mpf(b)) <= mp.mpf('1e-70')*max(1,abs(mp.mpf(b)))
        exact.append(b); expected.append(float(b))
    return {'names': NAMES, 'input_numerators': points, 'input_denominator': DENOMINATOR,
            'expected': expected, 'expected_120_digits': exact, 'batch_size': batch_size,
            'reference': 'Independent iterative parser and mpmath evaluation at 80 and 120 decimal digits',
            'source_bytes': len(source.encode()), 'source_function_occurrences': {
                f: source.count(f+'(') for f in ['exp','sin','cos']}}


def mathematica_syntax(source):
    result = []; stack = []; offset = 0; previous = None
    while offset < len(source):
        match = TOKEN.match(source, offset)
        if not match:
            assert not source[offset:].strip(); break
        token = match[1]; offset = match.end()
        if token == '(':
            function = previous in {'exp','sin','cos'}
            result.append('[' if function else '('); stack.append(']' if function else ')')
        elif token == ')': result.append(stack.pop())
        else: result.append({'exp':'Exp','sin':'Sin','cos':'Cos'}.get(token,token))
        previous = token
    assert not stack
    return ''.join(result)


def wolfram_operations(program):
    """Extend the polynomial counter with elementary calls and scalar division."""
    calls = Counter(); lines = []
    for line in program.splitlines():
        m = re.match(r'^(\d+\s+[RI]\d+ = )(.+)$', line)
        if m:
            rhs = m[2].strip()
            call = re.fullmatch(r'(Exp|Sin|Cos)\[\s*R\d+\s*\]',rhs)
            if call:
                calls[call[1].lower()] += 1; line = m[1]+'R0'
            elif re.fullmatch(r'(?:R\d+|1(?:\.0*)?)\s*/\s*R\d+',rhs):
                calls['div'] += 1; line = m[1]+'R0'
        lines.append(line)
    counted = common_polynomial_operations('\n'.join(lines))
    counted['arithmetic']['div'] = calls.pop('div',0)
    counted['arithmetic_operations'] = sum(counted['arithmetic'].values())
    counted['function_calls'] = dict(calls)
    counted['compile_print_sha256'] = hashlib.sha256(program.encode()).hexdigest()
    return counted


common.NAMES = NAMES
common.DENOMINATOR = DENOMINATOR
common_polynomial_operations = common.wolfram_operations
common.wolfram_operations = wolfram_operations
# Move Wolfram parsing into the generation phase, just like Symbolica and FORM.
parse = 'expression=ToExpression[StringReplace[Import["SOURCE","Text"],{"\\n"->"","\\r"->""}]];'
assert parse in common.WOLFRAM_NATIVE
common.WOLFRAM_NATIVE = common.WOLFRAM_NATIVE.replace(parse, '').replace(
    'If[Length[expression]!=reference["terms"],Exit[2]];', '').replace(
    'commands={}; compilerStart=Null; generationStart=SessionTime[];',
    'commands={}; compilerStart=Null; generationStart=SessionTime[];\n'+parse).replace(
    'Export["WORK/generation.done","done","Text"];',
    'If[Head[cf]===CompiledFunction,Apply[cf,Transpose[Take[points,4]]]];\n'
    'build=SessionTime[]-generationStart;\n'
    'Export["WORK/generation.done","done","Text"];')


def symbolica_route(route, source, reference, work, repeats):
    import numpy as np
    from symbolica import E,S
    text = source.read_text(); params = [S(n) for n in NAMES]
    points = np.asarray(reference['input_numerators'],dtype=np.float64)/DENOMINATOR
    batch = np.ascontiguousarray(points[np.arange(reference['batch_size'])%len(points)])
    common.gate(work,'generation'); start = time.perf_counter()
    expression = E(text)
    evaluator = expression.evaluator(params,iterations=10,n_cores=1,jit_compile=False,direct_translation=True)
    prepared = time.perf_counter()
    if route == 'symbolica_asm':
        executable = evaluator.compile('nested',str(work/'generated.cpp'),str(work/'generated.so'),
            'real',inline_asm='x64',optimization_level=3,compiler_path='g++',native=True,
            compiler_flags=['-fno-fast-math','-fno-unsafe-math-optimizations'])
    else:
        evaluator.jit_compile(True,direct_translation=True,optimization_level=3); executable=evaluator
    executable.evaluate(points[:4]); finished=time.perf_counter(); common.phase_done(work,'generation')
    instructions=evaluator.get_instructions(); counts=Counter(); calls=Counter()
    assert instructions.input_count == len(NAMES) and not instructions.sub_evaluators
    for inst in instructions.instructions:
        if inst[0] in {'add','mul'}: counts[inst[0]] += len(inst[2])-1
        elif inst[0] == 'pow':
            exponent=inst[3]; assert isinstance(exponent,int) and exponent != 0
            if exponent < 0: counts['div'] += 1; exponent=-exponent
            counts['mul'] += exponent.bit_length()+exponent.bit_count()-2
        elif inst[0] == 'fun':
            name=str(inst[2]).split('::')[-1]; assert name in {'exp','sin','cos'};calls[name]+=1
        else: raise ValueError(f'Unsupported instruction {inst}')
    verification=common.check_results(executable.evaluate(points).ravel(),reference['expected'])
    common.gate(work,'evaluation');samples=common.throughput(executable.evaluate,batch,repeats)
    common.phase_done(work,'evaluation')
    return {'version':version('symbolica'),'argument_count':len(NAMES),'kernel_callbacks':0,
        'optimization_seconds':prepared-start,'native_compilation_seconds':finished-prepared,
        'total_build_seconds':finished-start,'arithmetic':dict(counts),'arithmetic_operations':sum(counts.values()),
        'function_calls':dict(calls),'verification':verification,'evaluation_batch_seconds':samples,
        'evaluation_seconds_per_point':statistics.median(samples)/len(batch),
        'settings':{'iterations':10,'n_cores':1,'backend':route,'optimization_level':3,
                    'direct_translation':True,'fast_math':False}}


def cpp_operations(text):
    counts=Counter(); calls=Counter()
    def visit(node):
        if isinstance(node,ast.BinOp):
            op={ast.Add:'add',ast.Sub:'add',ast.Mult:'mul',ast.Div:'div'}.get(type(node.op));assert op
            counts[op]+=1
        elif isinstance(node,ast.Call):
            assert isinstance(node.func,ast.Name)
            if node.func.id == 'pow':
                exponent=ast.literal_eval(node.args[1]);assert isinstance(exponent,int) and exponent>=1
                counts['mul']+=exponent.bit_length()+exponent.bit_count()-2
                visit(node.args[0]);return
            assert node.func.id in {'exp','sin','cos'};calls[node.func.id]+=1
        elif isinstance(node,ast.UnaryOp) and isinstance(node.op,ast.USub) and not isinstance(node.operand,ast.Constant):
            counts['neg']+=1
        for child in ast.iter_child_nodes(node):visit(child)
    for statement in text.split(';'):
        if '=' in statement:visit(ast.parse(re.sub(r'\s+',' ',statement.split('=',1)[1]).strip(),mode='eval'))
    return {'arithmetic':dict(counts),'arithmetic_operations':sum(counts.values()),'function_calls':dict(calls)}


def form_route(route,source,reference,work,repeats):
    import numpy as np
    level='O2' if route.endswith('_o2') else 'O4'
    (work/'run.frm').write_text(f'''Off Statistics;
Symbols {','.join(NAMES)};
CFunctions exp,sin,cos;
.sort
#system kill -STOP $PPID
Local F =
#include {source}
;
.sort
ExtraSymbols,array,w;
Format C;
Format {level};
#optimize F
#write <generated.inc> "double w[1+`optimmaxvar_'];"
#write <generated.inc> "%O"
#write <generated.inc> "double result = %e",F
#write "FORM_OPERATIONS:`optimvalue_'"
.end
''')
    with (work/'engine.log').open('w') as log:
        process=subprocess.Popen(['form','-q','run.frm'],cwd=work,stdout=log,stderr=log)
        deadline=time.monotonic()+120
        while True:
            if process.poll() is not None:raise RuntimeError((work/'engine.log').read_text())
            if re.search(r'^State:\s+T',Path(f'/proc/{process.pid}/status').read_text(),re.M):break
            if time.monotonic()>deadline:raise TimeoutError('FORM startup')
            time.sleep(.005)
        points=np.asarray(reference['input_numerators'],dtype=np.float64)/DENOMINATOR
        batch=np.ascontiguousarray(points[np.arange(reference['batch_size'])%len(points)])
        common.gate(work,'generation');start=time.perf_counter();process.send_signal(signal.SIGCONT);process.wait()
    assert process.returncode==0,(work/'engine.log').read_text()
    generated=(work/'generated.inc').read_text()
    declarations='\n'.join(f'const double {name}=p[{i}];' for i,name in enumerate(NAMES))
    cpp=('#include <cmath>\n#include <cstddef>\nusing std::exp;using std::sin;using std::cos;\n'
         'static double nested_scalar(const double* p) {\n'+declarations+'\n'+generated+
         '\nreturn result;\n}\nextern "C" void nested_batch(const double* p,double* out,std::size_t n) {\n'
         f'for(std::size_t i=0;i<n;++i) out[i]=nested_scalar(p+{len(NAMES)}*i);\n}}\n')
    (work/'generated.cpp').write_text(cpp);prepared=time.perf_counter()
    command=['g++','-shared','-fPIC','-O3','-march=native','-std=c++17',str(work/'generated.cpp'),'-o',str(work/'generated.so')]
    compiler=subprocess.run(command,capture_output=True,text=True);(work/'compiler.log').write_text(compiler.stdout+compiler.stderr)
    assert compiler.returncode==0,compiler.stderr
    library=ctypes.CDLL(str(work/'generated.so'));function=library.nested_batch
    array=np.ctypeslib.ndpointer(dtype=np.float64,flags='C_CONTIGUOUS')
    function.argtypes=[array,array,ctypes.c_size_t];function.restype=None
    def evaluate(p):
        out=np.empty(len(p),dtype=np.float64);function(p,out,len(p));return out
    evaluate(points[:4]);finished=time.perf_counter();common.phase_done(work,'generation')
    verification=common.check_results(evaluate(points),reference['expected'])
    common.gate(work,'evaluation');samples=common.throughput(evaluate,batch,repeats);common.phase_done(work,'evaluation')
    return {'version':subprocess.check_output(['form','-v'],text=True).splitlines()[0],
        'argument_count':len(NAMES),'kernel_callbacks':0,'optimization_seconds':prepared-start,
        'native_compilation_seconds':finished-prepared,'total_build_seconds':finished-start,
        'compiler_commands':[command],'settings':{'form_optimization':level,'includes_required_expansion':True},
        'verification':verification,'evaluation_batch_seconds':samples,
        'evaluation_seconds_per_point':statistics.median(samples)/len(batch),**cpp_operations(generated)}


def worker(args):
    os.sched_setaffinity(0,{args.engine_cpu});work=args.work.resolve()
    reference=json.loads(args.reference.read_text());source=args.input.resolve()
    assert hashlib.sha256(source.read_bytes()).hexdigest()==reference['input_sha256']
    if args.worker.startswith('symbolica'):result=symbolica_route(args.worker,source,reference,work,args.repeats)
    elif args.worker.startswith('form'):result=form_route(args.worker,source,reference,work,args.repeats)
    else:
        translated=work/'input-mathematica.txt';translated.write_text(mathematica_syntax(source.read_text()))
        result=common.mathematica_route(args.worker,translated,args.reference.resolve(),reference,work,args.repeats)
    result.update(engine=args.worker,status='completed')
    assert result['argument_count']==7 and result['kernel_callbacks']==0
    (work/'result.json').write_text(json.dumps(result,indent=2)+'\n')


def main():
    parser=argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--input',type=Path,default=Path('benchmarks/large_nested.txt'))
    parser.add_argument('--output',type=Path,default=Path('scripts/benchmarks/large-nested-output-results.json'))
    parser.add_argument('--smoke',action='store_true');parser.add_argument('--runs',type=int,default=3)
    parser.add_argument('--routes',nargs='+',choices=ROUTES,default=ROUTES)
    parser.add_argument('--repeats',type=int,default=7);parser.add_argument('--batch-size',type=int,default=1024)
    parser.add_argument('--timeout',type=float,default=0,help='Optional build timeout in seconds; zero disables it')
    parser.add_argument('--min-available-mib',type=float,default=4096)
    parser.add_argument('--engine-cpu',type=int,default=0)
    parser.add_argument('--validation',type=Path,default=Path('scripts/benchmarks/large-nested-smoke-results.json'))
    parser.add_argument('--artifacts-dir',type=Path,default=Path('.quarto/large-nested-output'))
    parser.add_argument('--worker',choices=ROUTES);parser.add_argument('--reference',type=Path);parser.add_argument('--work',type=Path)
    args=parser.parse_args()
    if args.worker:worker(args);return
    assert args.runs>=1 and args.repeats>=1 and args.batch_size>=32
    digest=hashlib.sha256(Path(__file__).read_bytes()).hexdigest();helper_hash=hashlib.sha256(HELPER.read_bytes()).hexdigest()
    if not args.smoke:
        validation=json.loads(args.validation.read_text())
        assert validation['benchmark_source_sha256']==digest and validation['helper_source_sha256']==helper_hash
        assert all(any(r['engine']==engine and r['status']=='completed' for r in validation['results']) for engine in args.routes)
    affinity=sorted(os.sched_getaffinity(0));assert args.engine_cpu in affinity and len(affinity)>=2
    profiler=next(cpu for cpu in affinity if cpu!=args.engine_cpu);os.sched_setaffinity(0,{profiler})
    root=(args.artifacts_dir/('smoke-' if args.smoke else 'full-')/datetime.now(timezone.utc).strftime('%Y%m%dT%H%M%S%f')).resolve()
    root.mkdir(parents=True);source=root/'input.txt';source.write_text(SMOKE if args.smoke else args.input.read_text())
    print('Preparing independent high-precision reference…',flush=True)
    reference=reference_data(source.read_text(),args.batch_size);reference['input_sha256']=hashlib.sha256(source.read_bytes()).hexdigest()
    ref=root/'reference.json';ref.write_text(json.dumps(reference))
    data={'schema_version':1,'measured_at_utc':datetime.now(timezone.utc).isoformat(),'smoke':args.smoke,
        'cpu':next(line.split(':',1)[1].strip() for line in Path('/proc/cpuinfo').read_text().splitlines() if line.startswith('model name')),
        'platform':platform.platform(),'python':platform.python_version(),'input':reference,
        'benchmark_source_sha256':digest,'helper_source_sha256':helper_hash,'runs_per_engine':args.runs,
        'engine_cpu':args.engine_cpu,'profiler_cpu':profiler,'generation_timeout_seconds':args.timeout,
        'compiler_version':subprocess.check_output(['g++','--version'],text=True).splitlines()[0],
        'artifacts_dir':str(root),'results':[],
        'method':{'generation':'Includes parsing, required expansion, optimization, code generation, native compilation and first evaluation; engine startup/imports excluded',
                  'memory':'10 ms external worker/engine/compiler process-tree RSS samples; baseline before input parsing; profiler excluded',
                  'verification':'32 seeded mixed-sign dyadic points, independent 80/120 digit mpmath reference, relative tolerance 1e-9',
                  'evaluation':f'{args.batch_size}-point native batches, warmup then {args.repeats} samples; allocation included; one CPU'}}
    args.output.parent.mkdir(parents=True,exist_ok=True);failed=set()
    for run in range(1,args.runs+1):
        for route in (args.routes if run%2 else args.routes[::-1]):
            if route in failed:continue
            work=root/f'{run}-{route}';work.mkdir();print(f'Run {run}/{args.runs}: {route}',flush=True)
            command=[sys.executable,str(Path(__file__).resolve()),'--worker',route,'--input',str(source),
                     '--reference',str(ref),'--work',str(work),'--engine-cpu',str(args.engine_cpu),'--repeats',str(args.repeats)]
            memory=None;elapsed=0;status='worker_error'
            with (work/'worker.log').open('w') as log:
                process=subprocess.Popen(command,stdout=log,stderr=log,start_new_session=True)
                try:
                    status,elapsed,memory=common.measure_phase(process,work,'generation',args.timeout,args.min_available_mib)
                    if status=='completed':
                        evaluation_status,_,evaluation_memory=common.measure_phase(process,work,'evaluation',180,args.min_available_mib)
                        assert evaluation_status=='completed'
                        process.wait(timeout=180);assert process.returncode==0
                        result=json.loads((work/'result.json').read_text());result['evaluation_memory']=evaluation_memory
                    else:raise RuntimeError(status)
                except BaseException as error:
                    common.stop_tree(process.pid);process.wait()
                    if isinstance(error,(KeyboardInterrupt,SystemExit)):raise
                    logs='\n'.join((work/n).read_text(errors='replace') for n in ['worker.log','engine.log','compiler.log'] if (work/n).exists())
                    status='compiler_error' if 'internal compiler error' in logs else (status if status in {'generation_timeout','memory_limit'} else 'worker_error')
                    result={'engine':route,'status':status,'observed_generation_seconds':elapsed,'error':str(error),'logs':logs}
                result.update(run=run,generation_memory=memory,artifacts_dir=str(work));data['results'].append(result)
                args.output.write_text(json.dumps(data,indent=2)+'\n')
                print(f'  {result["status"]}: {result.get("total_build_seconds",elapsed):.3f} s',flush=True)
                if result['status']!='completed':
                    failed.add(route)
                    if args.smoke:raise RuntimeError('Smoke validation failed; see '+str(work))
    print('Saved '+str(args.output),flush=True)


if __name__=='__main__':main()