"""Verify production LSTP integration and probe force at real Amesim timestamps.

Dense trace sampling operates on an isolated runtime copy, without changing
the step schedule or time-event treatment. It never interpolates across events.
"""
from concurrent.futures import ThreadPoolExecutor
from dataclasses import replace
import argparse
import hashlib
import json
import math
import os
from pathlib import Path
import re
import shutil
import subprocess
import sys
from unittest.mock import patch

import numpy as np

ROOT=Path(__file__).resolve().parents[2]
sys.path.insert(0,str(ROOT))
from app.main import compile_system_xml_network
from app.simulation.backends import simulation_config
from app.simulation.native_codegen import build as builder,result_storage
from app.simulation.native_codegen.compiler import compile_native_program
from app.simulation.native_codegen.input import load_input
from app.simulation.native_codegen.runner import execute_native
from tests.manual import evaluate_mql8_correctness as evaluation,mql8_comparison as curves

PREVIOUS=ROOT/'test/mechanical-events-20260917/baselines'
DEFAULT=ROOT/'test/lstp-mainline-20260917'
save=evaluation.save


def model(profile):
    project=ROOT/'test/output-semantics-20260917/after'/profile/'platform.json'
    xml,doc=load_input(project)
    net=compile_system_xml_network(doc)
    return project,xml,doc,net,compile_native_program(net)


def build(program,out,native=None):
    with patch.object(builder,'CACHE',out/'cache'),patch.object(builder,'NATIVE',native or ROOT/'native'), \
         patch.object(builder,'ThreadPoolExecutor',lambda **kw:ThreadPoolExecutor(max_workers=1)):
        return builder.build_native(program)


def baseline(out):
    results={}
    for profile in ('full','noncyclic'):
        directory=out/'baseline'/profile
        if (directory/'summary.json').exists():
            results[profile]=json.loads((directory/'summary.json').read_bytes());continue
        directory.mkdir(parents=True,exist_ok=True)
        project,xml,doc,net,program=model(profile)
        old=json.loads((PREVIOUS/profile/'lstp/event-descriptors.json').read_bytes())
        # The sole generated-source addition is a constant descriptor table.
        original=re.sub(r'^const NativeContact model_contacts\[[^\n]+\n','',program.source,flags=re.M)
        assert hashlib.sha256(original.encode()).hexdigest()==old['originalModelSourceSha256']
        assert program.jacobian_structure==old['jacobianStructure']
        (directory/'platform.json').write_bytes(project.read_bytes());(directory/'platform.xml').write_bytes(xml)
        reference=directory/'amesim';reference.mkdir(exist_ok=True)
        for filename in ('test_mql_.var','test_mql_.results'):
            if not (reference/filename).exists():os.link(PREVIOUS/profile/'reference'/filename,reference/filename)
        audit=json.loads((PREVIOUS/profile/'audit/audit.json').read_bytes())
        print(profile,'production build',flush=True)
        compiled=build(program,out)
        try:
            with patch.object(result_storage,'RESULT_ROOT',out/'results'):
                r=execute_native(compiled,replace(simulation_config(doc.simulation),rtol=1e-8),.01,
                    run_dir=directory/'native',timeout=180)
            assert r['success'],r['message']
            summary=evaluation.compare(directory,json.loads(project.read_bytes()),net,audit,evaluation.PROFILES[profile])
            summary['nativeRun']={k:v for k,v in r.items() if k not in ('series','final','finalState')}
            with np.load(directory/'curves.npz') as current,np.load(PREVIOUS/profile/'lstp/curves.npz') as expected:
                summary['allComparedArraysEqualIsolatedLstp']=all(np.array_equal(current[k],expected[k],equal_nan=True) for k in current.files)
            summary['generatedRhsUnchanged']=True
            save(directory/'summary.json',summary);results[profile]=summary
            print(profile,'done',summary['allComparedArraysEqualIsolatedLstp'],r['contactEvents'],flush=True)
        finally:compiled.close()
    save(out/'baseline-summary.json',results)


def trace(out):
    project,xml,doc,net,program=model('full')
    audit=json.loads((PREVIOUS/'full/audit/audit.json').read_bytes());curves.configure(audit,net)
    ame=curves.read_ame(PREVIOUS/'full/reference')
    at=np.array(ame.times)
    # Query BEFORE and AFTER every forcing boundary, including the exact saved
    # Amesim row. These queries never become solver stop times.
    queries=set();points=[]
    boundaries=evaluation.event_times(json.loads(project.read_bytes()),50)
    for boundary in boundaries:
        if boundary<=0:continue
        ai=int(np.argmin(abs(at-boundary)))
        points.append(dict(boundary=boundary,ameIndex=ai,ameTime=float(at[ai])))
        queries.update((math.nextafter(boundary,-math.inf),boundary,float(at[ai])))
        for delta in (-1e-6,-1e-9,-1e-12,1e-14,1e-13,5e-13,1e-12,2e-12,3e-12,1e-11,1e-10,1e-9,1e-8,1e-7,1e-6,1e-5,1e-4,.001):
            queries.add(boundary+delta)
    queries=sorted(t for t in queries if 0<t<50)
    directory=out/'trace';directory.mkdir(exist_ok=True)
    save(directory/'queries.json',dict(points=points,times=queries,stateKeys=program.state_keys))
    runtime=directory/'native-source';shutil.copytree(ROOT/'native',runtime,dirs_exist_ok=True)
    source=runtime/'runtime/common.c';text=source.read_text(encoding='utf-8')
    hook='''
static void force_trace(double t,double end,int event,NativeDense dense,void *context) {
    static const double times[]={TIMES};
    static size_t cursor=0;
    while(cursor<sizeof(times)/sizeof(times[0]) && (event?times[cursor]<end:times[cursor]<=end)) {
        double at=times[cursor++],y[NSTATES];
        if(at<t || !dense(context,at,y)) continue;
        fprintf(stderr,"{\\"phase\\":\\"force-dense-trace\\",\\"time\\":%.17g,\\"state\\":[",at);
        for(int i=0;i<NSTATES;i++)fprintf(stderr,"%s%.17g",i?",":"",y[i]);
        fprintf(stderr,"]}\\n");
    }
}
'''.replace('TIMES',','.join(repr(t) for t in queries))
    text=text.replace('int native_accept(',hook+'\nint native_accept(',1)
    needle='    for (int i=0;i<count;i++) stop=fmin(stop,when[i]);'
    assert text.count(needle)==1
    text=text.replace(needle,needle+'\n    force_trace(t,stop,count,dense,context);',1)
    source.write_text(text,encoding='utf-8')
    print('dense trace build',flush=True);compiled=build(program,out,runtime)
    try:
        r=execute_native(compiled,replace(simulation_config(doc.simulation),rtol=1e-8),.01,
            record_samples=False,run_dir=directory/'run',timeout=180)
        assert r['success'],r['message']
        base=json.loads((out/'baseline/full/summary.json').read_bytes())['nativeRun']
        fields=('nfev','acceptedSteps','rejectedSteps','stateTransitions','solverStarts','njev','nlu','contactEvents')
        assert all(r[k]==base[k] for k in fields),'Trace changed integration work'
        states=[]
        for line in (directory/'run/worker.log').read_text(encoding='utf-8').splitlines():
            try:e=json.loads(line)
            except ValueError:continue
            if e.get('phase')=='force-dense-trace':states.append(e)
        assert [s['time'] for s in states]==queries
        payload='\n'.join(' '.join(format(v,'.17g') for v in (s['time'],*s['state'])) for s in states)+'\n'
        probes=subprocess.run([str(compiled.executable),'--probe'],input=payload,capture_output=True,text=True,check=True,timeout=45)
        output=[json.loads(line) for line in probes.stdout.splitlines()]
        keys=[v.key for v in program.variables]
        focused=[k for k in keys if any(word in k for word in ('amesim_mecmas21','amesim_lstp00a','amesim_ud00','amesim_forc'))]
        rows=[]
        for s,e in zip(states,output):
            assert e['success'];outputs=dict(zip(keys,e['outputs']))
            rows.append(dict(time=s['time'],values={k:outputs[k] for k in focused}))
        save(directory/'dense-outputs.json',rows)
        save(directory/'summary.json',dict(countersUnchanged=True,queries=len(rows),nativeRun={k:v for k,v in r.items() if k not in ('series','final','finalState')}))
        print('trace done',len(rows),flush=True)
    finally:compiled.close()


def rebase(out):
    """Short diagnostic continuation in a local clock, with the SAME saved
    boundary state and physics. This is not a production time-handling patch.
    """
    _,_,doc,net,program=model('full')
    points=json.loads((out/'trace/queries.json').read_bytes())['points']
    point=next(p for p in points if abs(p['boundary']-32.4)<1e-6)
    origin=point['boundary'];duration=point['ameTime']-origin
    summary=json.loads((out/'baseline/full/summary.json').read_bytes())
    state_file=out/'results'/summary['nativeRun']['resultStorage']['id']/'states.bin'
    assert not result_storage.scan_blocks(state_file)['corrupt']
    initial=None
    with state_file.open('rb') as stream:
        while header:=stream.read(result_storage._HEADER.size):
            magic,sequence,nrow,ncol,crc=result_storage._HEADER.unpack(header)
            assert magic==b'SIMBLK01' and 0<nrow<=1024 and ncol==len(program.state_keys)+1
            block=np.frombuffer(stream.read(nrow*ncol*8),dtype='<f8').reshape(ncol,nrow).T
            assert stream.read(8)==b'COMMIT01'
            matches=np.flatnonzero(block[:,0]==origin)
            if len(matches):initial=block[matches[-1],1:].copy()
    assert initial is not None
    original=program.source
    names=('model_init','model_eval','model_eval_jacobian','model_eval_jacobian_reuse',
           'model_friction_drives','model_property_temperatures','model_next_break')
    source=re.sub(r'\b('+ '|'.join(names)+r')\s*\(',lambda m:'absolute_'+m[1]+'(',original)
    source+='\nint model_init(double *y) {const double initial[NSTATES]={'+','.join(repr(float(v)) for v in initial)+'};memcpy(y,initial,sizeof(initial));return 1;}\n'
    for name,tail,args in [
        ('model_eval','double *dy,double *w','dy,w'),
        ('model_eval_jacobian','double *dy,double *w','dy,w'),
        ('model_eval_jacobian_reuse','double *dy,double *w,ModelJacobianWorkspace *workspace','dy,w,workspace'),
        ('model_friction_drives','double *drives','drives'),
        ('model_property_temperatures','NativePropertyTemperatures *temperatures','temperatures')]:
        source+=f'int {name}(double t,const double *y,{tail}) {{return absolute_{name}(t+{origin!r},y,{args});}}\n'
    # All STEP/UD00 signals are constant in this 2.2 ps continuation.
    source+='double model_next_break(double t,double end) {(void)t;return end;}\n'
    directory=out/'rebase';directory.mkdir(exist_ok=True)
    save(directory/'initial-state.json',dict(time=origin,stateKeys=program.state_keys,state=initial.tolist()))
    compiled=build(replace(program,source=source),out)
    try:
        records=[]
        for rtol in (1e-8,1e-10):
            config=replace(simulation_config(doc.simulation),t_start=0,t_stop=duration,max_step=duration/20,rtol=rtol)
            with patch.object(result_storage,'RESULT_ROOT',out/'results'):
                r=execute_native(compiled,config,duration/40,run_dir=directory/str(rtol),timeout=60)
            assert r['success'],r['message']
            record={k:v for k,v in r.items() if k not in ('series',)}
            record.update(origin=origin,duration=duration,rtol=rtol)
            records.append(record)
            print('rebase',rtol,r['final']['amesim_lstp00a_2.force'],r['solverControl'],flush=True)
        save(directory/'summary.json',records)
    finally:compiled.close()


def analyze(out):
    project,_,_,net,program=model('full')
    info=json.loads((out/'trace/queries.json').read_bytes())
    dense={row['time']:row['values'] for row in json.loads((out/'trace/dense-outputs.json').read_bytes())}
    descriptors=json.loads((PREVIOUS/'full/lstp/event-descriptors.json').read_bytes())['contacts']
    grid=np.load(out/'baseline/full/curves.npz')
    rows=[]
    for point in info['points']:
        boundary=point['boundary'];at=point['ameTime'];i=int(np.argmin(abs(grid['time']-boundary)))
        if not any(abs(boundary-t)<1e-6 for t in (21.6,32.4,43.2)):continue
        for contact in descriptors:
            name=contact['name'];_,v1,v2,gap,k,d,pdis,option=contact['values']
            keys=[program.state_keys[v1],program.state_keys[v2]]
            values={}
            for side in ('platform','amesim'):
                p=-float(grid[side+'|'+name+'.gap'][i])
                velocity=float(grid[side+'|'+keys[0]][i]-grid[side+'|'+keys[1]][i])
                elastic=k*p;damping=-math.expm1(-p/pdis)*d*velocity
                force=float(grid[side+'|'+name+'.force'][i])
                values[side]=dict(penetration=p,relativeVelocity=velocity,elastic=elastic,damping=damping,force=force,
                    reconstructedForce=elastic+damping,residual=force-(elastic+damping))
            error=values['platform']['force']-values['amesim']['force']
            same=dense[at][name+'.force'];same_error=same-values['amesim']['force']
            row=dict(component=name,nominalTime=float(grid['time'][i]),boundary=boundary,ameTime=at,
                elapsedAfterBoundary=at-boundary,**values,originalForceError=error,
                elasticError=values['platform']['elastic']-values['amesim']['elastic'],
                dampingError=values['platform']['damping']-values['amesim']['damping'],
                sameTimeForce=same,sameTimeForceError=same_error,
                errorReductionPercent=100*(1-abs(same_error)/abs(error)),
                timeUlp=math.ulp(boundary),forceChangePerTimeUlp=d*abs(dense[at]['amesim_mecmas21_10.a'])*math.ulp(boundary))
            rows.append(row)
    save(out/'force-error-analysis.json',rows)
    for row in rows:
        if row['component']=='amesim_lstp00a_2':print(json.dumps(row,ensure_ascii=False))
    # Report the absolute raw event peak separately from reference error.
    raw=json.loads((out/'baseline/full/native/result.json').read_bytes())['series']
    name='amesim_lstp00a_2';i=int(np.argmax(np.abs(raw[name+'.force'])))
    peak=dict(time=raw['time'][i],force=raw[name+'.force'][i],gap=raw[name+'.gap'][i],
        relativeVelocity=raw[name+'.port_1.v'][i]-raw[name+'.port_2.v'][i])
    peak['elastic']=-peak['gap']*1e11;peak['damping']=peak['force']-peak['elastic']
    keys=(name+'.force',name+'.gap',name+'.port_1.v',name+'.port_2.v',
          'amesim_mecmas21_10.x','amesim_mecmas21_10.v','amesim_ud00_2.out.signal')
    peak['adjacentSamples']=[dict(time=raw['time'][j],**{k:raw[k][j] for k in keys}) for j in (i-1,i,i+1)]
    save(out/'raw-force-peak.json',peak);print('raw peak',peak)
    grid.close()


def amesim_events(out):
    """Enable only Amesim's documented discontinuities printout on a copy."""
    directory=out/'amesim-event-output'
    evaluation.prepare_ame(ROOT/'tests/data/test_mql.ame',directory,50,.01,1e-8)
    sim=directory/'test_mql_.sim';lines=sim.read_text(encoding='ascii').splitlines()
    before=lines[:];options=lines[1].split()
    # Amesim 2404 scripting/python/amesim.py: ameputsimopt maps printDiscont
    # to simOptions[2]; every other solver and model setting stays unchanged.
    options[2]='1';lines[1]=' '.join(options);sim.write_text('\n'.join(lines)+'\n',encoding='ascii')
    save(directory/'print-option-change.json',dict(before=before,after=lines,source='Amesim 2404 ameputsimopt: simOptions[2]'))
    run=evaluation.run_ame(directory,Path('F:/AMESim2404/Amesim'))
    reference=curves.read_ame(directory)
    audit=json.loads((PREVIOUS/'full/audit/audit.json').read_bytes())
    aliases={x['component']:x['ameAlias'] for x in audit['mapping']}
    keys={'force':'f1@'+aliases['amesim_lstp00a_2'],
          'gap':'gap@'+aliases['amesim_lstp00a_2'],
          'massVelocity':'v1@'+aliases['amesim_mecmas21_10'],
          'branchVelocity':'v1@'+aliases['amesim_mecmas21_2'],
          'massPosition':'x1@'+aliases['amesim_mecmas21_10'],
          'drive':'output@'+aliases['amesim_ud00_2']}
    times=np.array(reference.times);data={k:np.array(reference.series(v)) for k,v in keys.items()}
    data['gap']*=.001
    selected=np.flatnonzero((times>32.39999)&(times<32.4004))
    rows=[dict(time=float(times[i]),**{k:float(v[i]) for k,v in data.items()}) for i in selected]
    j=int(np.argmax(abs(data['force'])))
    peak=dict(index=j,time=float(times[j]),**{k:float(v[j]) for k,v in data.items()})
    save(directory/'mechanical-events.json',dict(run=run,samples=len(times),around324=rows,maximumAbsoluteForce=peak))
    print('Amesim event output',len(times),'rows',rows,'peak',peak,flush=True)


def event_comparison(out):
    directory=out/'event-output-comparison';directory.mkdir(exist_ok=True)
    (directory/'native').mkdir(exist_ok=True);(directory/'amesim').mkdir(exist_ok=True)
    for source,target in [(out/'baseline/full/native/result.json',directory/'native/result.json'),
                          *[(out/'amesim-event-output'/name,directory/'amesim'/name) for name in ('test_mql_.var','test_mql_.results')]]:
        if not target.exists():os.link(source,target)
    project,_,_,net,_=model('full')
    audit=json.loads((PREVIOUS/'full/audit/audit.json').read_bytes())
    original_reference=curves.read_ame(directory/'amesim')
    times=np.array(original_reference.times);order=np.argsort(times,kind='stable')
    # Amesim writes its .8 s discontinuity after an already-written grid row
    # at .8000000000000009. Preserve every row/value and stable-sort by its
    # actual timestamp for this diagnostic matcher; retain the permutation.
    sorted_reference=replace(original_reference,times=tuple(times[order]),
        series_by_data_path={k:tuple(np.asarray(v)[order]) for k,v in original_reference.series_by_data_path.items()})
    save(directory/'ame-row-order.json',dict(originalTimeInversions=np.flatnonzero(np.diff(times)<0).tolist(),sortedToOriginal=order.tolist()))
    with patch.object(curves,'read_ame',return_value=sorted_reference):
        summary=evaluation.compare(directory,json.loads(project.read_bytes()),net,audit,evaluation.PROFILES['full'])
    # Adding event outputs must preserve every original Amesim saved sample.
    old=curves.read_ame(PREVIOUS/'full/reference');new=sorted_reference
    oldtimes=np.array(old.times);newtimes=np.array(new.times)
    indices=np.searchsorted(newtimes,oldtimes)
    assert np.array_equal(newtimes[indices],oldtimes)
    unchanged=all(np.array_equal(np.array(new.series(key))[indices],values) for key,values in old.series_by_data_path.items())
    assert unchanged,'Amesim event printout changed original results'
    summary['allOriginalAmesimRowsUnchanged']=unchanged
    save(directory/'summary.json',summary)
    print('event-output comparison',summary['groups']['force']['worstAbsolute'],
          'above5',summary['above5PercentCount'],'unpaired',summary['phaseUnpairedGridCount'],flush=True)


def main():
    parser=argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--output',type=Path,default=DEFAULT)
    parser.add_argument('--stage',choices=('baseline','trace','rebase','analyze','amesim-events','event-comparison','all'),default='all')
    args=parser.parse_args();out=args.output.resolve();out.mkdir(parents=True,exist_ok=True)
    if args.stage in ('baseline','all'):baseline(out)
    if args.stage in ('trace','all'):trace(out)
    if args.stage in ('rebase','all'):rebase(out)
    if args.stage in ('analyze','all'):analyze(out)
    if args.stage in ('amesim-events','all'):amesim_events(out)
    if args.stage in ('event-comparison','all'):event_comparison(out)


if __name__=='__main__':main()