"""Jacobian-local reuse: complete keys, isolation and uncached C parity."""
from pathlib import Path
import subprocess
import tempfile
import unittest

from app.simulation.native_codegen.build import toolchain, _command, COMPILER_FLAGS
from app.simulation.native_codegen.extended import compile_extended_program
from tests.test_native_catalog import Circuit
from tests.native_reference import reference_data, reference_network

ROOT=Path(__file__).resolve().parents[1]


def c_check(source,program=None):
    with tempfile.TemporaryDirectory() as temporary:
        directory=Path(temporary)
        if program:
            (directory/'model.h').write_text(program.header,encoding='utf-8')
            (directory/'model.c').write_text(program.source,encoding='utf-8')
        (directory/'test.c').write_text(source,encoding='utf-8')
        exe=directory/('test.exe' if __import__('os').name=='nt' else 'test')
        _command([toolchain()[0],*COMPILER_FLAGS,'-I',str(ROOT/'native/include'),'-I',str(directory),
                  str(directory/'test.c'),str(ROOT/'native/components/kernels.c'),'-lm','-o',str(exe)],log=[])
        result=subprocess.run([str(exe)],capture_output=True,text=True,timeout=60)
        if result.returncode: raise AssertionError(result.stdout+result.stderr)


class JacobianReuseTests(unittest.TestCase):
    def test_scalar_collision_capacity_keys_and_recording(self):
        c_check(r'''
#include "kernels.h"
#include <math.h>
#include <stdio.h>
#include <string.h>
#define CHECK(x) do {if(!(x)){fprintf(stderr,"line %d\n",__LINE__);return 1;}}while(0)
int main(void) {
 NativeJacobianScalarEntry entries[16]={0};
 NativeJacobianScalars cache={.entries=entries,.capacity=1,.recording=1};
 double key[10]={1,2,3,4,5,6,7,8,9,10},changed[10],value;
 native_jacobian_scalar_put(&cache,NATIVE_JACOBIAN_PH,1,key,10,11);
 CHECK(native_jacobian_scalar_get(&cache,NATIVE_JACOBIAN_PH,1,key,10,&value) && value==11);
 /* A one-slot cache guarantees collisions and saturation. */
 for(int i=0;i<10;i++) {
  memcpy(changed,key,sizeof(key));changed[i]=nextafter(changed[i],INFINITY);
  CHECK(!native_jacobian_scalar_get(&cache,NATIVE_JACOBIAN_PH,1,changed,10,&value));
  native_jacobian_scalar_put(&cache,NATIVE_JACOBIAN_PH,1,changed,10,12);
  CHECK(native_jacobian_scalar_get(&cache,NATIVE_JACOBIAN_PH,1,key,10,&value) && value==11);
 }
 CHECK(!native_jacobian_scalar_get(&cache,NATIVE_JACOBIAN_DENSITY,1,key,10,&value));
 CHECK(!native_jacobian_scalar_get(&cache,NATIVE_JACOBIAN_PH,0,key,10,&value));
 CHECK(!native_jacobian_scalar_get(&cache,NATIVE_JACOBIAN_PH,1,key,9,&value));
 memset(entries,0,sizeof(entries));cache.capacity=16;
 key[0]=0;native_jacobian_scalar_put(&cache,NATIVE_JACOBIAN_PH,1,key,10,1);
 key[0]=-0.;CHECK(!native_jacobian_scalar_get(&cache,NATIVE_JACOBIAN_PH,1,key,10,&value));
 native_jacobian_scalar_put(&cache,NATIVE_JACOBIAN_PH,1,key,10,NAN);
 CHECK(!native_jacobian_scalar_get(&cache,NATIVE_JACOBIAN_PH,1,key,10,&value));
 cache.recording=0;native_jacobian_scalar_put(&cache,NATIVE_JACOBIAN_PH,1,key,10,2);
 CHECK(!native_jacobian_scalar_get(&cache,NATIVE_JACOBIAN_PH,1,key,10,&value));
 key[0]=0;CHECK(native_jacobian_scalar_get(&cache,NATIVE_JACOBIAN_PH,1,key,10,&value) && value==1);
 cache.capacity=0;CHECK(!native_jacobian_scalar_get(&cache,NATIVE_JACOBIAN_PH,1,key,10,&value));
 cache.capacity=3;CHECK(!native_jacobian_scalar_get(&cache,NATIVE_JACOBIAN_PH,1,key,10,&value));
 return 0;
}
''')

    def test_scalar_reuse_replays_property_effects_even_with_full_local_cache(self):
        c_check(r'''
#include "kernels.h"
#include <math.h>
#include <stdio.h>
#include <string.h>
#define CHECK(x) do {if(!(x)){fprintf(stderr,"line %d\n",__LINE__);return 1;}}while(0)
int main(void) {
 NativeMedium m={1,2077.26439404998,5193.16098512495,293.15,0,1.96e-5,293.15,79.4};
 NativeJacobianScalarEntry entries[64]={0};
 NativeJacobianScalars memo={.entries=entries,.capacity=64,.recording=1};
 double initial[2];NativeGas gas;
 CHECK(native_medium_init(&m,2e5,300,.01,0,initial));
 CHECK(native_medium_gas_context(NULL,&m,initial[0],initial[1],.01,&gas));
 for(int iteration=0;iteration<4;iteration++) {
  NativePropertyState a[32],b[32];NativePropertyCache ca,cb;
  NativePropertyTemperatures ta={0},tb={0};
  native_properties_init(&ca,a,iteration%2?0:32);native_properties_init(&cb,b,iteration%2?0:32);
  ca.temperatures=&ta;cb.temperatures=&tb;cb.jacobian=&memo;
  double t1=native_temperature_ph_context(&ca,&m,2e5,gas.h),t2=native_temperature_ph_context(&cb,&m,2e5,gas.h);
  CHECK(!memcmp(&t1,&t2,sizeof(t1)));
  double r1=native_density_context(&ca,&m,2e5,t1),r2=native_density_context(&cb,&m,2e5,t2);
  CHECK(!memcmp(&r1,&r2,sizeof(r1)));
  double q1=native_pipe_flow_context(&ca,&m,2e5,1e5,t1,.001,1,.001,1);
  double q2=native_pipe_flow_context(&cb,&m,2e5,1e5,t2,.001,1,.001,1);
  CHECK(isfinite(q1) && !memcmp(&q1,&q2,sizeof(q1)));
  CHECK(ca.count==cb.count && ta.seen==tb.seen);
  for(int i=0;i<3;i++)CHECK(ta.minimum[i]==tb.minimum[i] && ta.maximum[i]==tb.maximum[i]);
  for(size_t i=0;i<ca.count;i++) {
   /* Compare every numerical field and validity flag; owner pointers differ. */
   a[i].temperatures=b[i].temperatures=NULL;a[i].jacobian=b[i].jacobian=NULL;
   CHECK(!memcmp(&a[i],&b[i],sizeof(a[i])));
  }
  memo.recording=0;
 }
 for(int i=0;i<NATIVE_JACOBIAN_SCALAR_KINDS;i++)CHECK(memo.reuses[i]>0);
 return 0;
}
''')

    def test_complete_bit_keys_immutable_baseline_failure_and_isolation(self):
        c_check(r'''
#include "kernels.h"
#include <math.h>
#include <stdio.h>
#include <string.h>
#define CHECK(x) do {if(!(x)){fprintf(stderr,"line %d\n",__LINE__);return 1;}}while(0)
static int calls;
static int kernel(NativePropertyCache *cache,const NativeMedium *m,double mass,double thermal,double volume,NativeGas *g) {
 if(cache)return 0;
 calls++;
 if(!isfinite(thermal))return 0;
 *g=(NativeGas){mass,thermal,volume,m->R,m->cp};return 1;
}
static int alternate(NativePropertyCache *c,const NativeMedium *m,double a,double b,double v,NativeGas *g) {
 return kernel(c,m,a,b,v,g);
}
int main(void) {
 NativeMedium m={1,2077,5193,293,0,1.96e-5,293,79.4};
 NativeJacobianGasMemo memo={0},other={0};NativeJacobianGasStats stats={0};NativeGas g,base;
 CHECK(native_jacobian_gas(&memo,kernel,&m,1,0,.01,&base,&stats));CHECK(calls==1);
 NativeMedium copy=m;
 CHECK(native_jacobian_gas(&memo,kernel,&copy,1,0,.01,&g,&stats));
 CHECK(calls==1 && stats.reuses==1 && !memcmp(&g,&base,sizeof(g)));
 double *fields[]={&copy.R,&copy.cp,&copy.Tref,&copy.slope,&copy.mu,&copy.muT,&copy.S};
 for(int i=0;i<7;i++) {
  copy=m;*fields[i]=nextafter(*fields[i],INFINITY);
  int before=calls;CHECK(native_jacobian_gas(&memo,kernel,&copy,1,0,.01,&g,&stats));CHECK(calls==before+1);
 }
 copy=m;copy.real_helium=0;
 int before=calls;CHECK(native_jacobian_gas(&memo,kernel,&copy,1,0,.01,&g,&stats));CHECK(calls==before+1);
 before=calls;CHECK(native_jacobian_gas(&memo,alternate,&m,1,0,.01,&g,&stats));CHECK(calls==before+1);
 before=calls;CHECK(native_jacobian_gas(&memo,kernel,&m,nextafter(1,2),0,.01,&g,&stats));CHECK(calls==before+1);
 before=calls;CHECK(native_jacobian_gas(&memo,kernel,&m,1,nextafter(0,1),.01,&g,&stats));CHECK(calls==before+1);
 before=calls;CHECK(native_jacobian_gas(&memo,kernel,&m,1,0,nextafter(.01,1),&g,&stats));CHECK(calls==before+1);
 before=calls;CHECK(native_jacobian_gas(&memo,kernel,&m,1,-0.,.01,&g,&stats));CHECK(calls==before+1 && signbit(g.T));
 before=calls;CHECK(!native_jacobian_gas(&memo,kernel,&m,1,NAN,.01,&g,&stats));CHECK(calls==before+1);
 /* Misses must not evict the original baseline; another workspace starts empty. */
 before=calls;CHECK(native_jacobian_gas(&memo,kernel,&m,1,0,.01,&g,&stats));CHECK(calls==before);
 CHECK(!memcmp(&g,&base,sizeof(g)));
 CHECK(!native_jacobian_gas(&other,kernel,&m,1,INFINITY,.01,&g,&stats));CHECK(!other.valid);
 CHECK(native_jacobian_gas(&other,kernel,&m,1,0,.01,&g,&stats));CHECK(other.valid);
 memset(&memo,0,sizeof(memo));before=calls;
 CHECK(native_jacobian_gas(&memo,kernel,&m,1,0,.01,&g,&stats));CHECK(calls==before+1);
 CHECK(stats.evaluations==(unsigned long)calls);
 return 0;
}
''')

    def check_program(self,program):
        c_check(r'''
#include "model.c"
#include <stdio.h>
#include <string.h>
#define CHECK(x) do {if(!(x)){fprintf(stderr,"line %d iteration %d\n",__LINE__,iteration);return 1;}}while(0)
int main(void) {
 int iteration=-1;double original[NSTATES],baseline[NSTATES],wo[NOUTPUTS];
 ModelJacobianWorkspace workspace={0};
 model_jacobian_begin(&workspace);
 CHECK(model_init(original));
 CHECK(model_eval_jacobian_reuse(.003,original,baseline,wo,&workspace));
 CHECK(workspace.stats.evaluations==MODEL_JACOBIAN_GAS_COUNT);
 for(iteration=0;iteration<2*NSTATES+4;iteration++) {
  double y[NSTATES],a[NSTATES],b[NSTATES],w[NOUTPUTS],v[NOUTPUTS];
  memcpy(y,original,sizeof(y));
  if(iteration<2*NSTATES)y[iteration/2]+=(iteration%2?1:-1)*1e-8*fmax(1,fabs(y[iteration/2]));
  if(iteration==2*NSTATES+1)y[0]=NAN;
  if(iteration==2*NSTATES+2)y[0]=INFINITY;
  int ok=model_eval_jacobian(.003,y,a,w);
  int cached=model_eval_jacobian_reuse(.003,y,b,v,&workspace);
  CHECK(ok==cached);
  if(ok)CHECK(!memcmp(a,b,sizeof(a)) && !memcmp(w,v,sizeof(w)));
 }
 CHECK(workspace.stats.reuses>0);
 /* Baseline remains usable after all perturbations and invalid trials. */
 unsigned long before=workspace.stats.evaluations;
 CHECK(model_eval_jacobian_reuse(.003,original,baseline,wo,&workspace));
 CHECK(workspace.stats.evaluations==before);
 return 0;
}
''',program)

    def test_thermal_storage_reverse_flow_and_moving_volume(self):
        b=Circuit();left=b.chamber('arbitrary_left',p0=8e4,T0=100,kth=0)
        pipe=b.add('amesim_pnl0002','arbitrary_pipe',p0=1e5,T0=1700,kth=0,diam=.001,le=1)
        right=b.chamber('arbitrary_right',p0=2e5,T0=300,kth=0)
        b.connect(left,'port_1',pipe,'port_1');b.connect(pipe,'port_2',right,'port_1')
        self.check_program(compile_extended_program(b.seal()))
        # Use catalog equations, including mechanical volume inputs, without
        # encoding any benchmark component names into the reuse implementation.
        case=next(c for c in reference_data()['cases'] if any(x['type']=='amesim_pnch012' for x in c['components']))
        self.check_program(compile_extended_program(reference_network(case)))

    def test_polytropic_and_ideal_medium_kernels(self):
        c_check(r'''
#include "kernels.h"
#include <math.h>
#include <stdio.h>
#include <string.h>
#define CHECK(x) do {if(!(x)){fprintf(stderr,"line %d\n",__LINE__);return 1;}}while(0)
int main(void) {
 NativeMedium m={0,287,1005,293,.01,1.8e-5,293,110.4};
 for(int helium=0;helium<2;helium++)for(int poly=0;poly<2;poly++) {
  m.real_helium=helium;if(helium){m.R=2077.26439404998;m.cp=5193.16098512495;}
  NativeCanonicalGas kernel=poly?native_polytropic_gas_context:native_medium_gas_context;
  NativeJacobianGasMemo memo={0};NativeJacobianGasStats stats={0};
  double temperatures[]={100,300,1700,25000,300};
  for(int i=0;i<5;i++) {
   double state[2];NativeGas a,b;
   CHECK(native_medium_init(&m,1e5,temperatures[i],.01,0,state));
   double thermal=poly?temperatures[i]:state[1];
   int ok=kernel(NULL,&m,state[0],thermal,.01,&a);
   int reused=native_jacobian_gas(&memo,kernel,&m,state[0],thermal,.01,&b,&stats);
   CHECK(ok==reused);if(ok)CHECK(!memcmp(&a,&b,sizeof(a)));
   reused=native_jacobian_gas(&memo,kernel,&m,state[0],thermal,.01,&b,&stats);
   CHECK(ok==reused);if(ok)CHECK(!memcmp(&a,&b,sizeof(a)));
  }
  CHECK(stats.reuses>0);
 }
 return 0;
}
''')


if __name__=='__main__': unittest.main()