"""Build/save a Symbolica evaluator and compare it with retained FORM binaries.

Run with a local Symbolica wheel plus NumPy. This diagnostic never recompiles
FORM and has no build timeout. Example (paths relative to the website checkout):

  .quarto/large-nested-fixed-env/bin/python scripts/benchmarks/compare-expanded-native.py \
    --provenance scripts/benchmarks/large-nested-fixed-provenance.json \
    --policy functions-first --cpu 8 --output .quarto/paired-functions-first

Use --smoke for the existing small validated fixtures, --dataset f13 for the
expanded F13 polynomial, and --load SAVED_DIRECTORY to reuse a saved evaluator.
Explicit orders/policies use the legacy tree path at zero search iterations so
that CSE/CPE still run. They are diagnostics, not automatic-search results.
"""
from __future__ import annotations

import argparse
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 statistics
import sys
import time

HERE = Path(__file__).resolve().parent
ROOT = HERE.parent.parent


def digest(path):
    return hashlib.sha256(Path(path).read_bytes()).hexdigest()


def write_json(path, value):
    temporary = path.with_suffix(path.suffix + '.tmp')
    temporary.write_text(json.dumps(value, indent=2) + '\n')
    temporary.replace(path)


def module_provenance(path):
    import symbolica
    modules = {Path(module.__file__).resolve() for name, module in sys.modules.items()
               if name.split('.')[0] == 'symbolica'
               and str(getattr(module, '__file__', '')).endswith('.so')}
    assert len(modules) == 1, modules
    native = modules.pop()
    result = {'module_path': str(native), 'native_module_sha256': digest(native),
              'package_version': version('symbolica')}
    if path:
        expected = json.loads(path.read_text())
        for field in ['native_module_sha256', 'package_version']:
            assert result[field] == expected[field], (field, result[field], expected[field])
        result['build_provenance'] = expected
        result['provenance_file'] = str(path.resolve())
        result['provenance_file_sha256'] = digest(path)
    return result


def defaults(dataset, smoke):
    if dataset == 'nested':
        primary = HERE / ('large-nested-smoke-results.json' if smoke else 'large-nested-form-original-results.json')
        forms = [('FORM O4', primary, 'form_cpp'),
                 ('FORM CSEGreedy', HERE / ('large-nested-form-csegreedy-smoke-results.json' if smoke
                                           else 'large-nested-form-csegreedy-results.json'), 'form_cpp_csegreedy')]
    else:
        primary = HERE / ('f13-subset-results.json' if smoke else 'f13-output-results.json')
        forms = [('FORM O4', primary, 'form_cpp'),
                 ('FORM O2', HERE / ('f13-form-o2-subset-results.json' if smoke else 'f13-form-o2-results.json'),
                  'form_cpp_o2')]
    return primary, forms


def check(values, expected):
    values = list(map(float, values))
    assert len(values) == len(expected)
    errors = []
    for actual, wanted in zip(values, expected):
        assert math.isfinite(actual)
        # Match the original fixture tolerance, while handling an exact zero.
        error = abs(actual - wanted) / abs(wanted) if wanted else abs(actual)
        assert error < 1e-9, (actual, wanted, error)
        errors.append(error)
    return {'verified': True, 'points': len(values), 'relative_tolerance': 1e-9,
            'max_relative_error': max(errors), 'values': values}


def count_instructions(evaluator, input_count):
    instructions = evaluator.get_instructions()
    assert instructions.input_count == input_count and not instructions.sub_evaluators
    arithmetic, functions = Counter(), Counter()
    for instruction in instructions.instructions:
        kind = instruction[0]
        if kind in {'add', 'mul'}:
            arithmetic[kind] += len(instruction[2]) - 1
        elif kind == 'pow':
            exponent = instruction[3]
            assert isinstance(exponent, int) and exponent
            if exponent < 0:
                arithmetic['div'] += 1
                exponent = -exponent
            arithmetic['mul'] += exponent.bit_length() + exponent.bit_count() - 2
        elif kind == 'fun':
            functions[str(instruction[2]).split('::')[-1]] += 1
        else:
            raise ValueError(f'Unsupported instruction: {instruction}')
    return {'arithmetic': dict(arithmetic), 'arithmetic_operations': sum(arithmetic.values()),
            'function_calls': dict(functions), 'instruction_count': len(instructions.instructions)}


def expand_all(expression):
    from symbolica import S, T
    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)
    return expression.expand(via_poly=True)


def select_order(expression, params, args, exp_wrapper=None):
    from symbolica import E, S
    argument = S('paired_function_argument_')
    def wrap_exp(value):
        return value.replace(S('exp')(argument), exp_wrapper(argument)) if exp_wrapper is not None else value
    if args.horner_order:
        order = [wrap_exp(E(item)) for item in args.horner_order]
        policy = 'explicit-order'
    elif args.horner_order_json:
        items = json.loads(args.horner_order_json.read_text())
        assert isinstance(items, list) and all(isinstance(item, str) for item in items)
        order = [wrap_exp(E(item)) for item in items]
        policy = 'explicit-order'
    elif args.policy == 'symbols-only':
        order, policy = params, args.policy
    elif args.policy == 'functions-first':
        # exp(argument) is represented internally as a power of Euler's constant,
        # so get_all_indeterminates alone does not enumerate it as a function.
        functions = {}
        for function in [exp_wrapper if exp_wrapper is not None else S('exp'), *S('sin', 'cos')]:
            for match in expression.match(function(argument)):
                value = function(match[argument])
                functions[str(value)] = value
        order = [functions[name] for name in sorted(functions)] + params
        policy = args.policy
    else:
        return None, {'policy': 'automatic', 'order': None,
                      'order_note': 'The Python evaluator API does not expose its selected automatic Horner order.'}
    assert len(set(map(str, order))) == len(order), 'Horner order must contain no duplicates'
    logical = [value.replace(exp_wrapper(argument), S('exp')(argument)) for value in order] if exp_wrapper is not None else order
    return order, {'policy': policy, 'order': list(map(str, logical)), 'applied_order': list(map(str, order)),
                   'order_note': 'Explicit diagnostic order; no order search. Coefficient grouping/content extraction then legacy-tree CSE/CPE.'}


def build_candidate(args, source, reference, points, provenance):
    from symbolica import E, S, FunctionDefinition
    started = time.perf_counter()
    expression = E(source.read_text())
    parsed = time.perf_counter()
    if args.dataset == 'nested':
        expression = expand_all(expression)
    expanded = time.perf_counter()
    params = [S(name) for name in reference['names']]
    explicit = args.policy != 'automatic' or args.horner_order or args.horner_order_json
    functions = []
    exp_wrapper = None
    if explicit and args.dataset == 'nested':
        # collect_horner cannot accept exp(argument), internally e^argument, as
        # an explicit indeterminate. Collect an opaque equivalent first, then
        # inline its exact definition while constructing the evaluator.
        exp_wrapper, argument, parameter = S('paired_exp_horner', 'paired_exp_argument_', 'paired_exp_parameter')
        expression = expression.replace(S('exp')(argument), exp_wrapper(argument))
        functions = [FunctionDefinition(exp_wrapper, [parameter], S('exp')(parameter), inlining='always')]
    order, ordering = select_order(expression, params, args, exp_wrapper)
    if exp_wrapper is not None:
        ordering['temporary_exp_representation'] = {'function': str(exp_wrapper),
            'definition': 'exp(parameter)', 'inlining': 'always',
            'reason': 'Python collect_horner rejects Euler-constant powers as explicit indeterminates.'}
    ordered = time.perf_counter()
    settings = {'iterations': args.iterations, 'n_cores': 1, 'cpe_iterations': None,
                'max_common_pair_cache_entries': args.pair_cache,
                'jit_compile': False, 'direct_translation': order is None}
    if order is not None:
        expression = expression.collect_horner(order).collect_by_coefficient().collect_num()
        settings['iterations'] = 0
    print('Horner policy:', ordering['policy'], 'order length:', None if order is None else len(order), flush=True)
    write_json(args.output / 'order.json', ordering)
    evaluator = expression.evaluator(params, functions=functions, **settings)
    prepared = time.perf_counter()
    evaluator.jit_compile(True, direct_translation=True, optimization_level=3)
    evaluator.evaluate(points[:4])
    finished = time.perf_counter()
    counts = count_instructions(evaluator, len(params))
    verification = check(evaluator.evaluate(points).ravel(), reference['expected'])
    saved = args.output / 'evaluator.bin'
    saved.write_bytes(bytes(evaluator.save()))
    # Serialization is part of this diagnostic's persistence, outside build time.
    restored = type(evaluator).load(saved.read_bytes())
    restore_verification = check(restored.evaluate(points).ravel(), reference['expected'])
    record = {'label': args.label, 'provenance': provenance, 'settings': settings,
              'jit_settings': {'direct_translation': True, 'optimization_level': 3},
              'input_form': 'expanded everywhere' if args.dataset == 'nested' else 'expanded polynomial',
              'ordering': ordering, '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': digest(saved),
              'verification': verification, 'serialization_verification': restore_verification,
              **counts}
    return restored, record


def load_candidate(args, reference, points, provenance):
    from symbolica import Evaluator
    build = json.loads((args.load / 'build.json').read_text())
    assert build['input_sha256'] == reference['input_sha256']
    record = build['candidate']
    assert record['provenance']['native_module_sha256'] == provenance['native_module_sha256']
    saved = args.load / 'evaluator.bin'
    assert digest(saved) == record['saved_evaluator_sha256']
    evaluator = Evaluator.load(saved.read_bytes())
    copied = args.output / 'evaluator.bin'
    copied.write_bytes(saved.read_bytes())
    record = dict(record, reused_from=str(args.load), saved_evaluator=str(copied))
    record['reload_verification'] = check(evaluator.evaluate(points).ravel(), reference['expected'])
    return evaluator, record


def load_form(label, path, engine, reference, points, dataset):
    import numpy as np
    data = json.loads(path.read_text())
    assert data['input']['input_sha256'] == reference['input_sha256'], (label, 'wrong input')
    assert data['input']['names'] == reference['names']
    assert data['input']['expected'] == reference['expected']
    rows = [row for row in data['results'] if row['engine'] == engine and row['status'] == 'completed']
    assert rows, (label, 'No completed FORM artifact')
    row = rows[0]
    binary = Path(row['artifacts_dir']) / 'generated.so'
    library = ctypes.CDLL(str(binary))
    function = getattr(library, 'nested_batch' if dataset == 'nested' else 'f13_batch')
    array = np.ctypeslib.ndpointer(dtype=np.float64, flags='C_CONTIGUOUS')
    function.argtypes = [array, array, ctypes.c_size_t]
    function.restype = None
    def evaluate(batch):
        output = np.empty(len(batch), dtype=np.float64)
        function(batch, output, len(batch))
        return output
    record = {'label': label, 'results_file': str(path.resolve()), 'results_file_sha256': digest(path),
              'original_engine': engine, 'original_run': row['run'],
              'binary': str(binary), 'binary_sha256': digest(binary),
              'generated_source_sha256': digest(binary.with_suffix('.cpp')),
              'original_settings': row.get('settings'),
              'verification': check(evaluate(points), reference['expected']),
              'reused_native_binary': True}
    counter_path = HERE / 'count-native-output.py'
    spec = importlib.util.spec_from_file_location('paired_form_counter', counter_path)
    counter = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(counter)
    generated = binary.with_suffix('.inc')
    record.update(counter.count_operations(generated.read_text()))
    record['operation_counter_source_sha256'] = digest(counter_path)
    record['generated_fragment_sha256'] = digest(generated)
    return evaluate, record


def balanced_orders(labels, rounds, seed):
    base = list(labels)
    random.Random(seed).shuffle(base)
    # Rotate each label through every position, then reverse the orientation.
    # Every complete block is balanced; partial blocks differ by at most one
    # occurrence per label/position, including the default twenty rounds.
    reverse = list(reversed(base))
    orders = [row[index:] + row[:index] for row in [base, reverse] for index in range(len(labels))]
    return [orders[index % len(orders)] for index in range(rounds)]


def percentile(values, fraction):
    values = sorted(values)
    position = (len(values) - 1) * fraction
    lower = int(position)
    return values[lower] + (values[min(lower + 1, len(values) - 1)] - values[lower]) * (position - lower)


def summarize(samples, labels, batch_size, seed):
    summary = {}
    for label in labels:
        values = [row['wall_ns'] / batch_size / 1e3 for row in samples if row['label'] == label]
        summary[label] = {'median_us_per_point': statistics.median(values),
                          'p10_us_per_point': percentile(values, .1),
                          'p90_us_per_point': percentile(values, .9), 'samples': len(values)}
    paired = []
    rounds = sorted({row['round'] for row in samples})
    by_round = {i: {row['label']: row['wall_ns'] for row in samples if row['round'] == i} for i in rounds}
    rng = random.Random(seed)
    for comparator in labels[1:]:
        ratios = [by_round[i][labels[0]] / by_round[i][comparator] for i in rounds]
        bootstrap = [statistics.median(rng.choices(ratios, k=len(ratios))) for _ in range(2000)]
        paired.append({'candidate': labels[0], 'comparator': comparator,
                       'median_runtime_ratio': statistics.median(ratios),
                       'median_ratio_bootstrap_95_interval': [percentile(bootstrap, .025), percentile(bootstrap, .975)],
                       'round_ratios': ratios,
                       'interpretation': 'Ratio below one means the candidate evaluates faster. CI resamples complete paired rounds.'})
    return {'engines': summary, 'paired_comparisons': paired}


def main():
    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument('--dataset', choices=['nested', 'f13'], default='nested')
    parser.add_argument('--smoke', action='store_true')
    parser.add_argument('--output', type=Path, required=True, help='New artifact directory; refuses to overwrite')
    parser.add_argument('--provenance', type=Path, help='Expected local wheel/source provenance JSON')
    parser.add_argument('--load', type=Path, help='Reuse evaluator.bin and build.json from a prior diagnostic')
    parser.add_argument('--label', default='Symbolica candidate')
    parser.add_argument('--iterations', type=int, default=10)
    parser.add_argument('--pair-cache', type=int, default=1000000)
    parser.add_argument('--policy', choices=['automatic', 'symbols-only', 'functions-first'], default='automatic')
    order = parser.add_mutually_exclusive_group()
    order.add_argument('--horner-order', nargs='+', help='Explicit Symbolica expressions in order, each shell-quoted')
    order.add_argument('--horner-order-json', type=Path, help='JSON array of Symbolica expression strings')
    parser.add_argument('--cpu', type=int, default=8)
    parser.add_argument('--rounds', type=int, default=20)
    parser.add_argument('--warmups', type=int, default=3)
    parser.add_argument('--seed', type=int, default=20261019)
    parser.add_argument('--batch-size', type=int)
    parser.add_argument('--build-only', action='store_true')
    parser.add_argument('--reference-results', type=Path, help='Override fixture result matrix, including its independent reference')
    parser.add_argument('--form', nargs=3, action='append', metavar=('LABEL', 'RESULTS_JSON', 'ENGINE'),
                        help='Override retained FORM baselines; repeat for each comparison')
    args = parser.parse_args()
    assert args.iterations >= 1 and args.pair_cache >= 1
    assert args.rounds >= 2 and args.warmups >= 1
    assert args.cpu in os.sched_getaffinity(0)
    assert args.provenance or args.load, 'Supply --provenance for a newly built local candidate'
    args.output = args.output.resolve()
    args.output.mkdir(parents=True, exist_ok=False)
    os.sched_setaffinity(0, {args.cpu})
    primary, forms = defaults(args.dataset, args.smoke)
    primary = args.reference_results or primary
    if args.form:
        forms = [(label, Path(path), engine) for label, path, engine in args.form]
    data = json.loads(primary.read_text())
    reference = data['input']
    source = Path(data['artifacts_dir']) / 'input.txt'
    assert digest(source) == reference['input_sha256']
    assert len(reference['input_numerators']) == len(reference['expected']) >= 32
    batch_size = args.batch_size or reference['batch_size']
    assert batch_size >= len(reference['expected'])
    import numpy as np
    points = np.ascontiguousarray(np.asarray(reference['input_numerators'], dtype=np.float64) / reference['input_denominator'])
    batch = np.ascontiguousarray(points[np.arange(batch_size) % len(points)])
    provenance = module_provenance(args.provenance)
    metadata = {'schema_version': 1, 'created_at_utc': datetime.now(timezone.utc).isoformat(),
                'dataset': args.dataset, 'smoke': args.smoke, 'input_sha256': reference['input_sha256'],
                'input_file': str(source), 'reference_results_file': str(primary.resolve()),
                'reference_results_sha256': digest(primary), 'reference': reference,
                'script_sha256': digest(__file__), 'python': sys.version, 'numpy_version': np.__version__,
                'platform': platform.platform(), 'cpu': args.cpu, 'affinity': sorted(os.sched_getaffinity(0)),
                'cpu_model': next(line.split(':', 1)[1].strip() for line in Path('/proc/cpuinfo').read_text().splitlines()
                                  if line.startswith('model name')),
                'thread_siblings': Path(f'/sys/devices/system/cpu/cpu{args.cpu}/topology/thread_siblings_list').read_text().strip(),
                'build_timeout_seconds': None}
    write_json(args.output / 'started.json', metadata)
    print('Building/loading candidate; no build timeout. Artifacts:', args.output, flush=True)
    evaluator, candidate = (load_candidate(args, reference, points, provenance) if args.load else
                            build_candidate(args, source, reference, points, provenance))
    write_json(args.output / 'build.json', dict(metadata, candidate=candidate))
    print('Verified and saved:', candidate['arithmetic_operations'], 'arithmetic operations;', candidate['function_calls'], flush=True)
    if args.build_only:
        return
    labels = [candidate['label']] + [label for label, _, _ in forms]
    assert len(labels) == len(set(labels)), 'Engine labels must be unique'
    evaluate = {labels[0]: lambda p: evaluator.evaluate(p).ravel()}
    comparators = []
    for label, path, engine in forms:
        print('Loading/checking retained binary:', label, flush=True)
        evaluate[label], comparator = load_form(label, path, engine, reference, points, args.dataset)
        comparators.append(comparator)
    for _ in range(args.warmups):
        for label in labels:
            evaluate[label](batch)
    orders = balanced_orders(labels, args.rounds, args.seed)
    result = dict(metadata, candidate=candidate, comparators=comparators, batch_size=batch_size,
                  rounds=args.rounds, warmups=args.warmups, seed=args.seed, orders=orders, samples=[],
                  method='Same-process same-core balanced interleaving; native scalar batches; output allocation included; all raw samples retained; no build timeout.')
    for round_index, current in enumerate(orders):
        for position, label in enumerate(current):
            cpu_start = time.process_time_ns()
            start = time.perf_counter_ns()
            output = evaluate[label](batch)
            elapsed = time.perf_counter_ns() - start
            cpu_elapsed = time.process_time_ns() - cpu_start
            assert np.isfinite(output).all()
            result['samples'].append({'round': round_index, 'position': position, 'label': label,
                                      'wall_ns': elapsed, 'process_cpu_ns': cpu_elapsed,
                                      'first_value': float(output[0]), 'last_value': float(output[-1])})
        write_json(args.output / 'results.json', result)
    result['final_verification'] = {label: check(evaluate[label](points), reference['expected']) for label in labels}
    result['summary'] = summarize(result['samples'], labels, batch_size, args.seed)
    result['completed_at_utc'] = datetime.now(timezone.utc).isoformat()
    result['completed'] = True
    write_json(args.output / 'results.json', result)
    print(json.dumps(result['summary'], indent=2), flush=True)


if __name__ == '__main__':
    main()