"""Exercise production reuse dispatch with public SUNDIALS types on either OS.

Only CVODE's weight/step getters are controlled; the production Jacobian,
colored/dense recovery and verification logic are compiled unchanged.
"""
from pathlib import Path
import shutil
import subprocess
import tempfile
import unittest

from app.simulation.native_codegen.build import toolchain, platform_build_inputs, link_library_arguments, _command

ROOT=Path(__file__).resolve().parents[1]
HEADER=r'''
#ifndef TEST_MODEL_H
#define TEST_MODEL_H
#include "kernels.h"
#define NSTATES 4
#define NOUTPUTS 4
#define MODEL_JACOBIAN_COLORED 1
#define MODEL_JACOBIAN_CANONICAL_RHS 1
#define MODEL_JACOBIAN_COLOR_COUNT 2
#define MODEL_JACOBIAN_NNZ 8
#define MODEL_JACOBIAN_GAS_REUSE 1
typedef struct {NativeJacobianGasStats stats;NativeJacobianScalars scalars;double baseline[4];int valid;} ModelJacobianWorkspace;
void model_jacobian_begin(ModelJacobianWorkspace *);
int model_eval_jacobian_reuse(double,const double *,double *,double *,ModelJacobianWorkspace *);
extern const int model_jacobian_column_color[4],model_jacobian_col_ptr[5],model_jacobian_row_index[8];
extern const double model_atol[4];
double model_next_break(double,double);
#endif
'''
HARNESS=r'''
#include <cvode/cvode.h>
#include <nvector/nvector_serial.h>
#include <sunmatrix/sunmatrix_dense.h>
#include <math.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
static int test_weights(void *solver,N_Vector weights){(void)solver;N_VConst(1e6,weights);return 0;}
static int test_step(void *solver,sunrealtype *step){(void)solver;*step=.01;return 0;}
#define CVodeGetErrWeights test_weights
#define CVodeGetCurrentStep test_step
#include "@RUNTIME@"
#define CHECK(x) do {if(!(x)){fprintf(stderr,"line %d\n",__LINE__);exit(1);}}while(0)
const int model_jacobian_column_color[4]={0,1,0,1},model_jacobian_col_ptr[5]={0,2,4,6,8};
const int model_jacobian_row_index[8]={0,1,0,1,2,3,2,3};
const double model_atol[4]={1e-9,1e-9,1e-9,1e-9};
static int cached_calls,uncached_calls,begins,reject_group,reject_individual,reject_base,corrupt;
static int polls,cancel_at;
static double base_state[4]={1,2,3,4};
static int changed(const double *y){int n=0;for(int i=0;i<4;i++)n+=y[i]!=base_state[i];return n;}
static void evaluate(const double *y,double *f) {
 for(int i=0;i<4;i+=2){f[i]=y[i]*y[i]+.5*y[i+1];f[i+1]=sin(y[i+1])+.25*y[i];}
}
int native_poll(NativeRun *r,double t){(void)r;(void)t;polls++;return !cancel_at || polls<cancel_at;}
int native_rhs(NativeRun *r,double t,const double *y,double *f){(void)t;r->nfev++;evaluate(y,f);return 1;}
int native_jacobian_rhs(NativeRun *r,double t,const double *y,double *f) {
 (void)t;r->nfev++;uncached_calls++;if(reject_individual && changed(y))return 0;evaluate(y,f);return 1;
}
void model_jacobian_begin(ModelJacobianWorkspace *w){memset(w,0,sizeof(*w));begins++;}
int model_eval_jacobian_reuse(double t,const double *y,double *f,double *outputs,ModelJacobianWorkspace *w) {
 (void)t;(void)outputs;cached_calls++;w->stats.evaluations++;w->scalars.evaluations[0]+=2;
 int n=changed(y);if((reject_base && !n) || (reject_group && n>1) || (reject_individual && n))return 0;
 if(!w->valid){memcpy(w->baseline,y,sizeof(w->baseline));w->valid=1;}
 evaluate(y,f);if(corrupt && n)f[0]+=1;return 1;
}
int native_fail(NativeRun *r,const char *a,const char *b,const char *c){(void)a;(void)b;(void)c;r->status=2;return 0;}
double model_next_break(double t,double end){(void)t;return end;}
int native_accept(NativeRun *r,double t,double next,const double *old,const double *trial,NativeDense dense,
                  void *context,double *at,double *accepted){(void)r;(void)t;(void)next;(void)old;(void)trial;(void)dense;(void)context;(void)at;(void)accepted;return -1;}
typedef struct {SUNContext sun;NativeRun run;CvContext context;N_Vector y,fy,tmp1,tmp2,tmp3;SUNMatrix matrix;} Fixture;
static void state(Fixture *f) {
 memcpy(N_VGetArrayPointer(f->y),base_state,sizeof(base_state));
 evaluate(base_state,N_VGetArrayPointer(f->fy));
 for(int i=0;i<4;i++)N_VGetArrayPointer(f->fy)[i]+=.125; /* ordinary fy determines increments */
}
static void init(Fixture *f) {
 memset(f,0,sizeof(*f));CHECK(!SUNContext_Create(SUN_COMM_NULL,&f->sun));
 f->y=N_VNew_Serial(4,f->sun);f->fy=N_VClone(f->y);f->tmp1=N_VClone(f->y);f->tmp2=N_VClone(f->y);f->tmp3=N_VClone(f->y);
 f->context=(CvContext){.run=&f->run,.colored=1,.weights=N_VClone(f->y)};f->run.jacobian_colored=1;
 f->context.jacobian_workspace=malloc(sizeof(*f->context.jacobian_workspace));CHECK(f->context.jacobian_workspace);
 f->matrix=SUNDenseMatrix(4,4,f->sun);state(f);
}
static int callback(Fixture *f){return cv_jacobian(.5,f->y,f->fy,f->matrix,&f->context,f->tmp1,f->tmp2,f->tmp3);}
static void matrix_correct(Fixture *f) {
 double increments[4],base[4],trial[4],value[4],fy[4];CHECK(!jac_increments(&f->context,f->y,f->fy,increments));
 evaluate(base_state,base);
 for(int j=0;j<4;j++) {
  memcpy(trial,base_state,sizeof(trial));trial[j]+=increments[j];evaluate(trial,value);
  double inverse=1.0/increments[j];
  for(int i=0;i<4;i++)CHECK(SM_ELEMENT_D(f->matrix,i,j)==inverse*(value[i]-base[i]));
 }
 CHECK(!memcmp(base_state,N_VGetArrayPointer(f->y),sizeof(base_state)));
 evaluate(base_state,fy);for(int i=0;i<4;i++)fy[i]+=.125;
 CHECK(!memcmp(fy,N_VGetArrayPointer(f->fy),sizeof(fy)));
}
int main(void) {
 Fixture f;init(&f);CHECK(valid_coloring());
 CHECK(!callback(&f));CHECK(cached_calls==3 && uncached_calls==0 && begins==1);matrix_correct(&f);
 CHECK(f.run.nfev==3 && f.run.jacobian_rhs==3 && f.run.jacobian_gas_evaluations==3 && f.run.jacobian_scalar_evaluations[0]==6);
 for(int i=0;i<4;i++){base_state[i]+=.125;}state(&f);
 CHECK(!callback(&f));CHECK(begins==2 && !memcmp(f.context.jacobian_workspace->baseline,base_state,sizeof(base_state)));matrix_correct(&f);
 /* Optional allocation failure retains ordinary canonical coloring. */
 ModelJacobianWorkspace *saved=f.context.jacobian_workspace;f.context.jacobian_workspace=NULL;
 int prior=cached_calls,prior_uncached=uncached_calls;
 CHECK(!callback(&f));CHECK(cached_calls==prior && uncached_calls-prior_uncached==3);matrix_correct(&f);
 f.context.jacobian_workspace=saved;
 int cached=cached_calls,uncached=uncached_calls;unsigned long before=f.run.nfev;
 reject_group=1;CHECK(!callback(&f));CHECK(cached_calls-cached==2 && uncached_calls-uncached==4 && f.run.nfev-before==6);
 CHECK(f.run.jacobian_fallbacks==1 && f.context.colored);matrix_correct(&f);reject_group=0;
 cached=cached_calls;uncached=uncached_calls;reject_individual=1;
 CHECK(callback(&f)>0);CHECK(cached_calls-cached==2 && uncached_calls-uncached==1);reject_individual=0;
 CHECK(f.run.jacobian_fallbacks==2);
 cached=cached_calls;uncached=uncached_calls;reject_base=1;
 CHECK(callback(&f)>0);CHECK(cached_calls-cached==1 && uncached_calls==uncached);reject_base=0;
 CHECK(f.run.jacobian_fallbacks==2);
 for(int first=1;first<=2;first++) {
  cached=cached_calls;uncached=uncached_calls;before=f.run.nfev;polls=0;cancel_at=first;
  CHECK(callback(&f)<0);CHECK(cached_calls-cached==first-1 && uncached_calls==uncached && f.run.nfev-before==(unsigned)(first-1));
  CHECK(f.run.jacobian_fallbacks==2);cancel_at=0;
 }
 f.context.reference=SUNDenseMatrix(4,4,f.sun);corrupt=1;cached=cached_calls;uncached=uncached_calls;
 CHECK(!callback(&f));CHECK(cached_calls-cached==3 && uncached_calls-uncached==4);
 CHECK(f.run.jacobian_checks==1 && f.run.jacobian_mismatches==1 && f.run.jacobian_fallbacks==3);
 CHECK(!f.context.colored && !f.run.jacobian_colored);matrix_correct(&f);
 cached=cached_calls;uncached=uncached_calls;int resets=begins;
 CHECK(!callback(&f));CHECK(cached_calls==cached && uncached_calls-uncached==5 && begins==resets);matrix_correct(&f);
 CHECK(f.run.nfev==f.run.jacobian_rhs);
 SUNMatDestroy(f.context.reference);SUNMatDestroy(f.matrix);
 free(f.context.jacobian_workspace);
 N_VDestroy(f.context.weights);N_VDestroy(f.y);N_VDestroy(f.fy);N_VDestroy(f.tmp1);N_VDestroy(f.tmp2);N_VDestroy(f.tmp3);SUNContext_Free(&f.sun);
 puts("reuse runtime cases passed");return 0;
}
'''


class JacobianReuseRuntimeTests(unittest.TestCase):
    def test_reset_reference_fallback_failure_and_cancellation(self):
        compiler,sundials,_=toolchain()
        flags,libraries,dlls,executable=platform_build_inputs(sundials)
        with tempfile.TemporaryDirectory() as directory:
            directory=Path(directory)
            (directory/'model.h').write_text(HEADER,encoding='utf-8')
            source=HARNESS.replace('@RUNTIME@',(ROOT/'native/runtime/cvode_solver.c').as_posix())
            (directory/'test.c').write_text(source,encoding='utf-8')
            target=directory/executable
            _command([compiler,*flags,'-I',str(directory),'-I',str(ROOT/'native/include'),'-I',str(sundials/'include'),
                      str(directory/'test.c'),*link_library_arguments(libraries),'-lm','-o',str(target)],log=[])
            for dll in dlls: shutil.copy2(dll,directory/dll.name)
            result=subprocess.run([str(target)],capture_output=True,text=True,timeout=30)
            self.assertEqual(result.returncode,0,result.stdout+result.stderr)
            self.assertIn('reuse runtime cases passed',result.stdout)


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