"""Isolated Jacobian-only experiment. Never installs a production build.

prepare builds a reference worker and an opt-in local-probe worker. The latter
keeps the original model_eval_internal verbatim and adds a canonical clone.
Only compiler-owned, acyclic schedule operations are eligible for skipping.
Exact live property-context and pipe-cache guards preserve graph-external
cache semantics. An audit compares every probe's complete dy/w to the original.
"""
from pathlib import Path
from concurrent.futures import ThreadPoolExecutor
import argparse, hashlib, json, os, re, shutil, subprocess, sys, time

ROOT=Path(__file__).resolve().parents[2]
sys.path.insert(0,str(ROOT))
from app.main import compile_system_xml_network
from app.simulation.native_codegen.input import load_input
from app.simulation.native_codegen.compiler import compile_native_program
from app.simulation.native_codegen.jacobian import StateDependencies
from app.simulation.native_codegen.schedule import EvaluationSchedule
from app.simulation.native_codegen import build as builder

OUT=ROOT/'test/local-probe-20260917'
TEMPLATE=Path(__file__).with_name('local_probe_support.c')
KERNELS=[('properties','native_medium_gas'),('properties','native_temperature_ph'),
 ('properties','native_density'),('properties','native_viscosity'),('properties','state_valve'),
 ('properties','property_pt'),('properties','local_isentropic'),
 ('pipe','native_pipe_resistance'),('pipe','native_pipe_flow_context'),
 ('pipe','native_pipe_diagnostics_context'),('orifice','native_medium_orifice_context')]

def write(path,value):path.write_text(json.dumps(value,ensure_ascii=False,indent=2)+'\n',encoding='utf-8')
def replace(s,old,new):
    assert s.count(old)==1,(old,s.count(old));return s.replace(old,new)
def function_span(s,name):
    m=re.search(r'(?m)^(?:static\s+|NATIVE_COMPONENT_INTERNAL\s+)?(?:int|double|void|NativePropertyState\s*\*)\s*'+re.escape(name)+r'\([^;{}]*\)\s*\{',s)
    assert m,name
    # Generated/native reviewed functions have no braces in string literals in
    # the functions selected below. Mask comments and strings for brace matching.
    masked=re.sub(r'/\*[\s\S]*?\*/|//[^\n]*|"(?:\\.|[^"\\])*"',lambda x:' '*len(x[0]),s)
    b=masked.index('{',m.start());depth=1;e=b+1
    while depth:depth+=(masked[e]=='{')-(masked[e]=='}');e+=1
    return m.start(),b,e

def dependency_masks(d):
    keys=set(d.seeds)|set(d.inputs)|set().union(*d.inputs.values())
    unknown=keys-set(d.seeds)-set(d.inputs);full=(1<<d.state_count)-1
    masks={k:full if k in unknown else d.seeds.get(k,0) for k in keys}
    for _ in range(len(keys)+1):
        changed=False
        for k,inputs in d.inputs.items():
            value=masks[k]
            for ref in inputs:value|=masks[ref]
            if value!=masks[k]:masks[k]=value;changed=True
        if not changed:return masks,unknown
    raise AssertionError('dependency closure did not converge')

def capture(model):
    saved={};old_build=StateDependencies.build;old_emit=EvaluationSchedule.emit
    def dep(self):saved['dependencies']=self;return old_build(self)
    def emit(self):
        saved['schedule']=self;result=old_emit(self);saved['schedule_lines']=result[1];return result
    StateDependencies.build=dep;EvaluationSchedule.emit=emit
    try:
        _,doc=load_input(model);network=compile_system_xml_network(doc);program=compile_native_program(network)
    finally:StateDependencies.build=old_build;EvaluationSchedule.emit=old_emit
    return program,saved

def generate(program,saved):
    schedule=saved['schedule'];d=saved['dependencies'];structure=d.build()
    if not structure.enabled or not program.manifest()['jacobianStructure']['reuse']['enabled']:
        raise ValueError('Experiment requires eligible canonical Jacobian reuse')
    if any(b.cyclic for b in schedule.blocks):
        raise ValueError('Experimental local schedule rejects SCCs; use the unmodified worker')
    if structure.color_count>63:raise ValueError('Experimental group mask supports at most 63 groups')
    order=[b.members[0] for b in schedule.blocks];ops=[schedule.computations[i] for i in order]
    masks,unknown=dependency_masks(d);full=(1<<d.state_count)-1
    opmask=[0]*len(ops)
    for i,op in enumerate(ops):
        for ref in op.inputs:opmask[i]|=masks.get(ref,full)
        if any(not re.fullmatch(r'(?:p|h|q|w|fb)\[\d+\]',v) for v in op.outputs):
            raise ValueError('Unsupported local operation outputs')
    regions=[];plans=[];groups=[]
    for color in range(structure.color_count):
        states=[j for j,c in enumerate(structure.colors) if c==color]
        mask=sum(1<<j for j in states);affected=[bool(v&mask) for v in opmask]
        plan=[-1]*len(ops);i=0
        while i<len(ops):
            if affected[i]:i+=1;continue
            start=i
            while i<len(ops) and not affected[i]:i+=1
            region=(start,i)
            if region not in regions:regions.append(region)
            plan[start]=regions.index(region)
        plans.append(plan)
        groups.append(dict(color=color,stateIndices=states,states=[program.state_keys[i] for i in states],
            affectedOperations=[order[i] for i,v in enumerate(affected) if v],
            regions=[list(regions[x]) for x in plan if x>=0]))
    boundaries=sorted({x for pair in regions for x in pair})
    # Pure assignments cannot mutate either context. Share their checkpoint
    # slot, retaining a new checkpoint after every unreviewed/native call.
    pure_calls={'if','for','sizeof','fmax','fmin','fabs','sqrt','copysign','pow'}
    versions=[0]
    for op in ops:
        calls=set(re.findall(r'\b([A-Za-z_]\w*)\s*\(', '\n'.join(op.code)))
        versions.append(versions[-1]+bool(calls-pure_calls))
    checkpoints=sorted({versions[x] for x in boundaries})
    boundary={x:checkpoints.index(versions[x]) for x in boundaries}
    source=program.source
    dims={name:int(re.search(pattern,source)[1]) for name,pattern in {
        'NP':r'double p\[(\d+)\]', 'NQ':r'h\[(\d+)\]', 'NFB':r'double fb\[(\d+)\]',
        'NPC':r'NativePipeCache pipe_cache\[(\d+)\]', 'NPS':r'NativePropertyState property_states\[(\d+)\]'}.items()}
    outputs=[];offsets=[0]
    arraynames=['p','h','q','w','fb']
    for op in ops:
        for target in op.outputs:
            a,n=re.fullmatch(r'(p|h|q|w|fb)\[(\d+)\]',target).groups();outputs.append((arraynames.index(a),int(n)))
        offsets.append(len(outputs))
    macros={'NC':structure.color_count,'NO':len(ops),'NR':len(regions),'NB':len(checkpoints),'NK':len(KERNELS),**dims}
    header='''#ifndef LOCAL_PROBE_EXPERIMENT_H
#define LOCAL_PROBE_EXPERIMENT_H
#include "model.h"
#include <stdint.h>
'''+''.join(f'#define LP_{k} {v}\n' for k,v in macros.items())+'''
extern int lp_color,lp_active,lp_capture,lp_observe;
extern uint64_t lp_mask;
extern const int lp_plan[LP_NC][LP_NO],lp_end[LP_NR];
void lp_start(void);void lp_finish(void);void lp_begin(double,const double*);void lp_ready(int);
void lp_snapshot(int,NativePropertyCache*,NativePipeCache*);
void lp_save(double*,double*,double*,double*,double*);
int lp_reuse(int,NativePropertyCache*,NativePipeCache*,double*,double*,double*,double*,double*);
#if LP_OBSERVE
void lp_operation(int);
#else
#define lp_operation(position) ((void)(position))
#endif
void lp_kernel(int);int lp_valid(double,const double*);
void lp_note_eval(int);
void lp_compare(int,const double*,const double*,int,const double*,const double*);
void lp_matrix(double,const double*,const double*);void lp_event(double,const double*);
uint64_t lp_tick(void);void lp_jac_time(uint64_t);void lp_newton(long,long);
int lp_eval(double,const double*,double*,double*,ModelJacobianWorkspace*);
#endif
'''
    declarations='const int lp_end[LP_NR]={'+','.join(str(b) for a,b in regions)+'};\n'
    declarations+='const int lp_plan[LP_NC][LP_NO]={'+','.join('{'+','.join(map(str,p))+'}' for p in plans)+'};\n'
    declarations+='static const int lp_before[LP_NR]={'+','.join(str(boundary[a]) for a,b in regions)+'};\n'
    declarations+='static const int lp_after[LP_NR]={'+','.join(str(boundary[b]) for a,b in regions)+'};\n'
    declarations+='static const int lp_first_output[LP_NR]={'+','.join(str(offsets[a]) for a,b in regions)+'};\n'
    declarations+='static const int lp_last_output[LP_NR]={'+','.join(str(offsets[b]) for a,b in regions)+'};\n'
    declarations+='static const int lp_start_op[LP_NR]={'+','.join(str(a) for a,b in regions)+'};\n'
    declarations+='static const int lp_output_map[][2]={'+','.join('{'+str(a)+','+str(i)+'}' for a,i in outputs)+'};\n'
    declarations+='static const int lp_operation_id[LP_NO]={'+','.join(map(str,order))+'};\n'
    declarations+='static const int lp_colors[NSTATES]={'+','.join(map(str,structure.colors))+'};\n'
    declarations+='static const char *lp_kernel_names[LP_NK]={'+','.join(json.dumps(f) for _,f in KERNELS)+'};\n'
    # Keep the original entire model evaluator byte-for-byte. Add an opt-in clone
    # whose only numerical statement change is replacement of the schedule.
    a,b,e=function_span(source,'model_eval_internal');original=source[a:e]
    schedule_text='\n'.join(saved['schedule_lines']);assert original.count(schedule_text)==1
    capture_lines=[]
    for pos,op in enumerate(ops):
        if pos in boundary:capture_lines.append(f'lp_snapshot({boundary[pos]},properties,pipe_cache);')
        capture_lines.extend(op.code)
    if len(ops) in boundary:capture_lines.append(f'lp_snapshot({boundary[len(ops)]},properties,pipe_cache);')
    capture_lines.append('lp_save(p,h,q,w,fb);')
    local=['if(lp_capture){',*capture_lines,'}else{','for(int pos=0;pos<LP_NO;){',
        'int region=lp_plan[lp_color][pos];',
        'if(region>=0 && lp_reuse(region,properties,pipe_cache,p,h,q,w,fb)){pos=lp_end[region];continue;}',
        'lp_operation(pos);','switch(pos){']
    for pos,op in enumerate(ops):local += [f'case {pos}:{{',*op.code,'break;}']
    local += ['default:return 0;}','pos++;','}}']
    clone=original.replace('model_eval_internal(', 'model_eval_local_internal(',1).replace(schedule_text,'\n'.join(local))
    wrapper='''
int lp_eval(double t,const double *y,double *dy,double *w,ModelJacobianWorkspace *workspace){
    static ModelJacobianWorkspace shadow;
    int active=lp_color+1;lp_active=active;
    int eligible=lp_mask && workspace && (lp_color<0 || (lp_mask&(UINT64_C(1)<<lp_color)));
    if(lp_color<0){lp_begin(t,y);if(lp_observe)model_jacobian_begin(&shadow);}
    int valid=lp_color<0 || lp_valid(t,y);
    lp_capture=lp_color<0;
    lp_note_eval(eligible && valid);
    int result=eligible && valid ? model_eval_local_internal(t,y,dy,w,1,NULL,NULL,workspace) : model_eval_jacobian_reuse(t,y,dy,w,workspace);
    if(eligible && valid && workspace)workspace->scalars.recording=0;
    if(lp_color<0)lp_ready(eligible && result);
    if(lp_observe){
        double expected_dy[NSTATES],expected_w[NOUTPUTS];lp_active=-1;
        int expected=model_eval_jacobian_reuse(t,y,expected_dy,expected_w,&shadow);
        lp_compare(result,dy,w,expected,expected_dy,expected_w);
    }
    lp_active=-1;return result;
}
'''
    local_source=source+'\n'+clone+'\n'+wrapper
    # The reference original is never replaced in the timing worker. Observation
    # counts reference operations analytically in support code, not by modifying it.
    assert original in local_source
    metadata=dict(stateKeys=program.state_keys,groups=groups,order=order,regions=regions,boundaries=boundaries,contextCheckpointSlots=boundary,
        operations=[dict(id=order[i],key=op.key,inputs=sorted(op.inputs),outputs=list(op.outputs),stateIndices=[j for j in range(d.state_count) if opmask[i]>>j&1]) for i,op in enumerate(ops)],
        unknownDependencyLeaves=sorted(unknown),macros=macros,originalEvaluatorSha256=hashlib.sha256(original.encode()).hexdigest())
    return '#include "local_probe.h"\n'+local_source,header,declarations,metadata

def kernel_counter(s,fn,index):
    a,b,e=function_span(s,fn)
    return s[:b+1]+f'if(lp_active>=0)lp_kernel({index});'+s[b+1:]

def prepare(model,observe=False):
    OUT.mkdir(parents=True,exist_ok=True);started=time.perf_counter();program,saved=capture(model)
    local,header,tables,meta=generate(program,saved)
    write(OUT/'plan.json',meta);(OUT/'original-model.c').write_text(program.source,encoding='utf-8')
    work=OUT/('audit' if observe else 'worker');work.mkdir(exist_ok=True)
    (work/'model.h').write_text(program.header,encoding='utf-8');(work/'local_probe.h').write_text(header,encoding='utf-8')
    sources={p.relative_to(builder.NATIVE).as_posix():p.read_text(encoding='utf-8') for p in builder._runtime_sources(program)}
    source_hashes={key:hashlib.sha256(value.encode()).hexdigest() for key,value in sources.items()}
    sources['model.c']=local
    common=sources['runtime/common.c']
    common=replace(common,'int ok=r->options.bdf ? native_bdf(r) : native_rk45(r);','lp_start();int ok=r->options.bdf ? native_bdf(r) : native_rk45(r);')
    common=replace(common,'r->solve_cpu_seconds=native_cpu_time()-r->cpu_start;','r->solve_cpu_seconds=native_cpu_time()-r->cpu_start;lp_finish();')
    common=replace(common,'r->events++;','lp_event(stop,accepted_state);r->events++;')
    sources['runtime/common.c']=common
    cv=sources['runtime/cvode_solver.c']
    cv=replace(cv,'model_eval_jacobian_reuse(t,N_VGetArrayPointer(y),N_VGetArrayPointer(f),outputs,workspace)',
        'lp_eval(t,N_VGetArrayPointer(y),N_VGetArrayPointer(f),outputs,workspace)')
    cv=replace(cv,'for (int color=0;color<MODEL_JACOBIAN_COLOR_COUNT;color++) {\n        memcpy(test,state,NSTATES*sizeof(double));',
        'for (int color=0;color<MODEL_JACOBIAN_COLOR_COUNT;color++) {\n        lp_color=color;memcpy(test,state,NSTATES*sizeof(double));')
    a,b,e=function_span(cv,'cv_jacobian');sig=cv[a:b].strip()
    impl=cv[a:e].replace('cv_jacobian(', 'cv_jacobian_original(',1)
    wrapper=sig+'{lp_color=-1;uint64_t start=lp_tick();int result=cv_jacobian_original(t,y,fy,matrix,user,tmp1,tmp2,tmp3);lp_jac_time(start);if(!result)lp_matrix(t,N_VGetArrayPointer(y),SUNDenseMatrix_Data(matrix));return result;}'
    cv=cv[:a]+impl+'\n'+wrapper+cv[e:]
    cv=replace(cv,'long int value=0;', 'long int value=0,nni=0,ncf=0;CVodeGetNumNonlinSolvIters(solver,&nni);CVodeGetNumNonlinSolvConvFails(solver,&ncf);lp_newton(nni,ncf);')
    sources['runtime/cvode_solver.c']=cv
    if observe:
        for i,(module,fn) in enumerate(KERNELS):
            key=f'components/modules/{module}.c';sources[key]=kernel_counter(sources[key],fn,i)
    sources={key:'#include "local_probe.h"\n'+src for key,src in sources.items()}
    sources['local_probe_support.c']=TEMPLATE.read_text(encoding='utf-8').replace('/* GENERATED_TABLES */',tables)
    cc,sun,_=builder.toolchain();flags,libs,dlls,exe=builder.platform_build_inputs(sun);flags += [f'-DLP_OBSERVE={int(observe)}']
    generated=time.perf_counter()
    def compile_one(item):
        i,(key,source)=item;path=work/Path(key).name;path.write_text(source,encoding='utf-8',newline='\n');obj=work/f'unit-{i}.o';log=[]
        builder._command([cc,*flags,'-I',str(work),'-I',str(builder.NATIVE/'include'),'-I',str(sun/'include'),'-c',str(path),'-o',str(obj)],log=log,timeout=180)
        return obj,log
    with ThreadPoolExecutor(max_workers=4) as pool:objects=list(pool.map(compile_one,enumerate(sources.items())))
    log=[];builder._command([cc,*flags,*[str(x) for x,_ in objects],*builder.link_library_arguments(libs),'-lm','-o',str(work/exe)],log=log)
    for dll in dlls:shutil.copyfile(dll,work/dll.name)
    (work/'build.log').write_text('\n'.join([*sum([x for _,x in objects],[]),*log]),encoding='utf-8')
    write(work/'build-metadata.json',dict(commit=subprocess.check_output(['git','rev-parse','HEAD'],cwd=ROOT,text=True).strip(),
        inputSha256=hashlib.sha256(model.read_bytes()).hexdigest(),sourceHashes=source_hashes,
        codegenSeconds=generated-started,compileSeconds=time.perf_counter()-generated,totalSeconds=time.perf_counter()-started))
    print('BUILT',work.name,'seconds',time.perf_counter()-started,'regions',len(meta['regions']),'boundaries',len(meta['boundaries']),flush=True)

def run(label,mask,observe=False,stop=10):
    work=OUT/label;work.mkdir(exist_ok=True);exe=OUT/('audit' if observe else 'worker')/'model.exe'
    env=os.environ.copy();env['LOCAL_PROBE_MASK']=str(mask)
    args=[str(exe),'--method','BDF','--start','0','--stop',str(stop),'--sample-step','.01','--max-step','1e30','--rtol','1e-8','--timeout','300',
        '--sample-file',str(work/'states.bin'),'--output-block-file',str(work/'outputs.bin'),'--output',str(work/'result.json')]
    start=time.perf_counter()
    with (work/'stderr.log').open('wb') as f:p=subprocess.run(args,cwd=work,env=env,stdout=subprocess.PIPE,stderr=f,timeout=330,creationflags=subprocess.CREATE_NO_WINDOW)
    elapsed=time.perf_counter()-start
    if p.returncode:raise RuntimeError((label,p.returncode,(work/'stderr.log').read_text(encoding='utf-8')[-5000:]))
    r=json.loads((work/'result.json').read_text(encoding='utf-8'));diag=json.loads((work/'probe.json').read_text())
    record={k:r[k] for k in ('success','finalState','final','propertyWarnings','acceptedSteps','rejectedSteps','stateTransitions','solverStarts','nfev','njev','nlu','solveSeconds','solveCpuSeconds')}
    for name in ('states','outputs','events','jacobians'):
        path=work/f'{name}.bin'
        if path.exists():
            with path.open('rb') as f:record[name+'Sha256']=hashlib.file_digest(f,'sha256').hexdigest()
            record[name+'Bytes']=path.stat().st_size
    record.update(processSeconds=elapsed,mask=mask,diagnostic=diag)
    write(work/'measurement.json',record)
    print('RUN',label,{k:record[k] for k in ('solveSeconds','solveCpuSeconds','processSeconds','acceptedSteps','nfev','njev','nlu')},'newton',diag['newtonIterations'],'audit',diag['auditComparisons'],diag['auditDifferences'],flush=True)

def cold_runs(model):
    for name,mask in [('reference',0),('all',0x7ffffff)]:
        start=time.perf_counter();prepare(model,False);prepared=time.perf_counter()-start
        run('cold-'+name,mask);elapsed=time.perf_counter()-start
        record=json.loads((OUT/('cold-'+name)/'measurement.json').read_text(encoding='utf-8'))
        write(OUT/('cold-'+name+'-pipeline.json'),dict(mode=name,prepareSeconds=prepared,
            workerSeconds=record['processSeconds'],pipelineSeconds=elapsed))

def guard_test():
    work=OUT/'audit';cc,sun,_=builder.toolchain();flags,_,_,_=builder.platform_build_inputs(sun);log=[]
    executable=work/'guard-test.exe'
    builder._command([cc,*flags,'-DLP_OBSERVE=1','-I',str(work),'-I',str(builder.NATIVE/'include'),
        str(Path(__file__).with_name('local_probe_guard_test.c')),'-lm','-o',str(executable)],log=log)
    target=OUT/'guard-selftest';target.mkdir(exist_ok=True);env=os.environ.copy();env['LOCAL_PROBE_MASK']=str(1<<26)
    subprocess.run([str(executable)],cwd=target,env=env,check=True,creationflags=subprocess.CREATE_NO_WINDOW)

if __name__=='__main__':
    p=argparse.ArgumentParser();p.add_argument('action',choices=['prepare','run','cold','guard']);p.add_argument('--audit',action='store_true')
    p.add_argument('--model',type=Path,default=ROOT/'tests/data/test-mql-8-corrected.json');p.add_argument('--label',default='run');p.add_argument('--mask',type=lambda v:int(v,0),default=0);p.add_argument('--stop',type=float,default=10)
    a=p.parse_args()
    if a.action=='prepare':prepare(a.model,a.audit)
    elif a.action=='cold':cold_runs(a.model)
    elif a.action=='guard':guard_test()
    else:run(a.label,a.mask,a.audit,a.stop)