SystemSimulationApp/tests/manual/mechanical_event_variant.py

"""Isolated mechanical-event builds; never edits production numerics or cache.

Augment the generated program with local contact descriptors and an optional
read-only release-drive evaluator. The ordinary RHS, Jacobian, property reuse,
state layout, output layout, and component formulas stay byte-for-byte intact.
"""
from dataclasses import replace
import hashlib
import json
from pathlib import Path
import re
import shutil

from app.simulation.components.amesim.semantics import contact_stiffness

HERE = Path(__file__).resolve().parent
MODES = ('off', 'mass', 'lstp', 'all')


def replace_once(text, old, new):
    if text.count(old) != 1:
        raise ValueError(f'Experiment patch no longer matches the reviewed runtime: {old[:100]}')
    return text.replace(old, new, 1)


def prepare_runtime(original, destination):
    original, destination = Path(original), Path(destination)
    shutil.copytree(original, destination, dirs_exist_ok=True)
    shutil.copyfile(HERE/'mechanical_events_runtime.c', destination/'runtime/experimental_events.c')
    common = destination/'runtime/common.c'
    source = common.read_text(encoding='utf-8')
    source = replace_once(source, 'int native_accept(', '#include "experimental_events.c"\n\nint native_accept(')
    previous_size=('2*(NSTOPS+NFRICTIONS+NCONTACTS+1)' if '2*(NSTOPS+NFRICTIONS+NCONTACTS+1)' in source
                   else '2*(NSTOPS+NFRICTIONS+1)')
    event_size=previous_size[:-2]+'NEXPERIMENT_CONTACTS+NEXPERIMENT_RELEASES+1)'
    source = source.replace(previous_size,event_size)
    for field in ('bounds','restitution','thresholds'):
        declaration=field+'['+event_size+']'
        if declaration+'={0}' not in source:
            source=replace_once(source,declaration,declaration+'={0}')
    source = replace_once(source, '    double stop=next;', '''    if (!experiment_candidates(r,t,next,old,trial,dense,context,when,indices,friction,&count)) return -1;
    double stop=next;''')
    if 'if(friction[i]>=0) continue;' in source:
        source = replace_once(source, 'if(friction[i]>=0) continue;', 'if(friction[i]!=-1) continue;')
    elif 'if(friction[i]!=-1) continue;' not in source:
        raise ValueError('Unknown hard-stop event dispatch')
    source = source.replace('if(friction[i]<0 &&', 'if(friction[i]==-1 &&')
    source = replace_once(source, '        if (!native_append(r,stop,accepted_state)) return -1;', '''        experiment_commit(r,stop,accepted_state,when,indices,friction,count);
        if (!native_append(r,stop,accepted_state)) return -1;''')
    common.write_text(source, encoding='utf-8')
    header = destination/'include/runtime.h'
    source = header.read_text(encoding='utf-8')
    source = replace_once(source, '    NativePropertyWarning property_warnings[6];', '''    double experiment_last[2*NEXPERIMENT_CONTACTS+NEXPERIMENT_RELEASES+1];
    int experiment_direction[2*NEXPERIMENT_CONTACTS+NEXPERIMENT_RELEASES+1];
    int experiment_pending[2*NEXPERIMENT_CONTACTS+NEXPERIMENT_RELEASES+1];
    unsigned long experiment_checks, experiment_dense, experiment_roots, experiment_rhs;
    unsigned long experiment_mass, experiment_lstp, experiment_clipping, experiment_release;
    NativePropertyWarning property_warnings[6];''')
    header.write_text(source, encoding='utf-8')
    main = destination/'runtime/main.c'
    source = main.read_text(encoding='utf-8')
    source = replace_once(source, '    fprintf(f,",\\\"propertyWarnings\\\":"); native_property_warnings_json(r,f);', '''    fprintf(f,",\\\"experimentalEvents\\\":{\\\"checks\\\":%lu,\\\"denseCalls\\\":%lu,\\\"rootIterations\\\":%lu,\\\"releaseRhsCalls\\\":%lu,\\\"massContactTransitions\\\":%lu,\\\"lstpContactTransitions\\\":%lu,\\\"forceClipTransitions\\\":%lu,\\\"massReleaseTransitions\\\":%lu}",
        r->experiment_checks,r->experiment_dense,r->experiment_roots,r->experiment_rhs,
        r->experiment_mass,r->experiment_lstp,r->experiment_clipping,r->experiment_release);
    fprintf(f,",\\\"propertyWarnings\\\":"); native_property_warnings_json(r,f);''')
    main.write_text(source, encoding='utf-8')


def event_program(program, network, mode):
    if mode not in MODES:
        raise ValueError(mode)
    # Reproduce opt-in experiments after production LSTP integration too:
    # descriptors remain available, but only the selected experimental events
    # run. This also keeps the tests portable when the frozen snapshot is absent.
    program=replace(program,header=re.sub(r'#define NCONTACTS \d+','#define NCONTACTS 0',program.header))
    # The off program uses the unmodified runtime and generated program.
    if mode == 'off':
        return program, dict(mode=mode, contacts=[], releases=[])
    slots = {v.key: i for i, v in enumerate(program.variables)}
    assignments = {}
    for w, y in re.findall(r'w\[(\d+)\]\s*=\s*y\[(\d+)\]\s*;', program.source):
        if int(w) in assignments and assignments[int(w)] != int(y):
            raise ValueError('Ambiguous kinematic state alias')
        assignments[int(w)] = int(y)

    def velocity(key):
        vi = assignments[slots[key+'.v']]
        if assignments[slots[key+'.x']] != vi+1:
            raise ValueError('Experiment requires the reviewed contiguous v/x state layout')
        return vi

    contacts = []
    for c in network.components.values():
        if c.model_type == 'amesim_lstp00a' and mode in ('lstp', 'all'):
            contacts.append(dict(name=c.name, values=(3, velocity(c.name+'.port_1'), velocity(c.name+'.port_2'),
                c.gap0, contact_stiffness(c), c.rcont, c.Pdis, int(c.discContactOption))))
        if c.model_type == 'amesim_mecmas21' and int(c.stoptype) == 2 and mode in ('mass', 'all'):
            for kind, side in ((1, 'min'), (2, 'max')):
                contacts.append(dict(name=c.name+'.'+side, values=(kind, velocity(c.name), -1,
                    getattr(c,'x'+side), getattr(c,'Kb'+side), getattr(c,'Db'+side),
                    getattr(c,'Pd'+side), int(c.discContactOption))))
    stop_count = int(re.search(r'#define NSTOPS (\d+)', program.header)[1])
    releases = list(range(stop_count)) if mode in ('mass', 'all') else []
    header = f'''
#define NEXPERIMENT_CONTACTS {len(contacts)}
#define NEXPERIMENT_RELEASES {len(releases)}
typedef struct {{ int kind, v1, v2; double boundary, stiffness, damping, pdis; int signed_force; }} ExperimentContact;
extern const ExperimentContact experiment_contacts[{max(1,len(contacts))}];
int model_experiment_release_drives(double t,const double *y,double *drives);
'''
    data = ','.join('{'+','.join(str(v) if isinstance(v,int) else repr(float(v)) for v in c['values'])+'}' for c in contacts)
    additions = f'\nconst ExperimentContact experiment_contacts[{max(1,len(contacts))}] = {{{data or "{0,0,0,0,0,0,0,0}"}}};\n'
    if releases:
        start = program.source.index('static int model_eval_internal(')
        body_start = program.source.index('{', start)
        depth, end = 1, body_start+1
        while depth:
            depth += (program.source[end] == '{') - (program.source[end] == '}')
            end += 1
        function = program.source[start:end]
        function = replace_once(function, 'model_eval_internal(', 'model_experiment_release_internal(')
        pos = function.index(') {')
        function = function[:pos] + ',double *release_drives' + function[pos:]
        found = []
        def capture(match):
            vi = int(re.search(r'&dy\[(\d+)\]', match[0])[1])
            index = len(found)
            found.append(vi)
            return f'release_drives[{index}]=dy[{vi}];'
        function = re.sub(r'native_stop_motion\([^;]+\);', capture, function)
        if len(found) != stop_count:
            raise ValueError('Release evaluator does not match generated stop count')
        extended = 'ModelJacobianWorkspace' in function
        args = '0,NULL,NULL,NULL,drives' if extended else 'NULL,drives'
        additions += function + '\nint model_experiment_release_drives(double t,const double *y,double *drives) {double dy[NSTATES],w[NOUTPUTS];return model_experiment_release_internal(t,y,dy,w,'+args+');}\n'
    else:
        additions += 'int model_experiment_release_drives(double t,const double *y,double *drives) {(void)t;(void)y;(void)drives;return 1;}\n'
    # Add declarations before the include guard closes, preserving all generated
    # RHS and Jacobian source bytes as an exact prefix.
    at = program.header.rfind('#endif')
    modified = replace(program, source=program.source+additions, header=program.header[:at]+header+program.header[at:])
    assert modified.source.startswith(program.source)
    return modified, dict(mode=mode, contacts=contacts, releases=releases,
        originalModelSourceSha256=hashlib.sha256(program.source.encode()).hexdigest(),
        jacobianStructure=program.jacobian_structure)