/* CVODE retains its Jacobian refresh policy and dense matrix/LU solver.
 * A compiler-proven sparsity pattern groups independent finite differences;
 * unsupported models and patterns without grouping benefit retain CVODE's
 * default callback.
 */
#include "runtime.h"
#include <cvode/cvode.h>
#include <cvode/cvode_ls.h>
#include <nvector/nvector_serial.h>
#include <sunmatrix/sunmatrix_dense.h>
#include <sunlinsol/sunlinsol_dense.h>
#include <math.h>
#include <string.h>
#include <sundials/sundials_math.h>
#ifndef MODEL_JACOBIAN_COLORED
#define MODEL_JACOBIAN_COLORED 0
#endif

typedef struct { void *solver; N_Vector scratch; } CvDense;
static int cv_dense(void *context, double t, double *out) {
    CvDense *d=context;
    if (CVodeGetDky(d->solver,t,0,d->scratch) < 0) return 0;
    memcpy(out,N_VGetArrayPointer(d->scratch),NSTATES*sizeof(double));
    return 1;
}
typedef struct {
    NativeRun *run;
    void *solver;
    N_Vector weights;
    SUNMatrix reference;
    int colored;
} CvContext;
static int cv_rhs(sunrealtype t, N_Vector y, N_Vector dy, void *context) {
    NativeRun *r=((CvContext *)context)->run;
    if (!native_poll(r,r->final_time)) return -1;
    return native_rhs(r,t,N_VGetArrayPointer(y),N_VGetArrayPointer(dy)) ? 0 : 1;
}

#if MODEL_JACOBIAN_COLORED
/* Check the generated CSC bounds and coloring once, before installing it. */
static int valid_coloring(void) {
    if (MODEL_JACOBIAN_COLOR_COUNT<=0 || MODEL_JACOBIAN_COLOR_COUNT>=NSTATES ||
        model_jacobian_col_ptr[0]!=0 || model_jacobian_col_ptr[NSTATES]!=MODEL_JACOBIAN_NNZ) return 0;
    for (int j=0;j<NSTATES;j++) {
        if (model_jacobian_column_color[j]<0 || model_jacobian_column_color[j]>=MODEL_JACOBIAN_COLOR_COUNT ||
            model_jacobian_col_ptr[j]<0 || model_jacobian_col_ptr[j+1]<model_jacobian_col_ptr[j] ||
            model_jacobian_col_ptr[j+1]>MODEL_JACOBIAN_NNZ) return 0;
        int previous=-1;
        for (int k=model_jacobian_col_ptr[j];k<model_jacobian_col_ptr[j+1];k++) {
            int row=model_jacobian_row_index[k];
            if (row<=previous || row>=NSTATES) return 0;
            previous=row;
        }
    }
    for (int color=0;color<MODEL_JACOBIAN_COLOR_COUNT;color++) {
        unsigned char used[NSTATES]={0};
        for (int j=0;j<NSTATES;j++) if (model_jacobian_column_color[j]==color) {
            for (int k=model_jacobian_col_ptr[j];k<model_jacobian_col_ptr[j+1];k++) {
                int row=model_jacobian_row_index[k];
                if (used[row]) return 0;
                used[row]=1;
            }
        }
    }
    return 1;
}

static int jac_rhs(CvContext *context, sunrealtype t, N_Vector y, N_Vector f) {
    if (!native_poll(context->run,context->run->final_time)) return -1;
    context->run->jacobian_rhs++;
#if defined(MODEL_JACOBIAN_CANONICAL_RHS) && MODEL_JACOBIAN_CANONICAL_RHS
    return native_jacobian_rhs(context->run,t,N_VGetArrayPointer(y),N_VGetArrayPointer(f)) ? 0 : 1;
#else
    return native_rhs(context->run,t,N_VGetArrayPointer(y),N_VGetArrayPointer(f)) ? 0 : 1;
#endif
}

/* Same perturbation and reciprocal-multiply order as SUNDIALS 7.4 dense DQ.
 * No constraints are set by this runtime. If that changes, carry their vector
 * into this context and apply CVODE's sign rule before enabling coloring. */
static int jac_increments(CvContext *context, N_Vector y, N_Vector fy, double *increments) {
    sunrealtype step;
    if (CVodeGetErrWeights(context->solver,context->weights)<0 ||
        CVodeGetCurrentStep(context->solver,&step)<0) return -1;
    double norm=N_VWrmsNorm(fy,context->weights);
    double minimum=norm!=0 ? 1000.0*fabs(step)*SUN_UNIT_ROUNDOFF*NSTATES*norm : 1.0;
    double *state=N_VGetArrayPointer(y), *weight=N_VGetArrayPointer(context->weights);
    double square_root=sqrt(SUN_UNIT_ROUNDOFF);
    for (int j=0;j<NSTATES;j++) {
        increments[j]=fmax(square_root*fabs(state[j]),minimum/weight[j]);
        if (!(increments[j]>0) || !isfinite(increments[j])) return 1;
    }
    return 0;
}

static int dense_difference(CvContext *context, sunrealtype t, N_Vector y, N_Vector fy,
                            SUNMatrix matrix, N_Vector trial, N_Vector ftrial, const double *increments) {
    double *state=N_VGetArrayPointer(y), *test=N_VGetArrayPointer(trial);
    double *base=N_VGetArrayPointer(fy), *value=N_VGetArrayPointer(ftrial);
    memcpy(test,state,NSTATES*sizeof(double));
    for (int j=0;j<NSTATES;j++) {
        test[j]=state[j]+increments[j];
        int flag=jac_rhs(context,t,trial,ftrial);
        test[j]=state[j];
        if (flag) return flag;
        double inverse=1.0/increments[j];
        double *column=SUNDenseMatrix_Column(matrix,j);
        for (int i=0;i<NSTATES;i++) column[i]=inverse*(value[i]-base[i]);
    }
    return 0;
}

static int colored_difference(CvContext *context, sunrealtype t, N_Vector y, N_Vector fy,
                              SUNMatrix matrix, N_Vector trial, N_Vector ftrial, const double *increments) {
    double *state=N_VGetArrayPointer(y), *test=N_VGetArrayPointer(trial);
    double *base=N_VGetArrayPointer(fy), *value=N_VGetArrayPointer(ftrial);
    if (SUNMatZero(matrix)) return -1;
    for (int color=0;color<MODEL_JACOBIAN_COLOR_COUNT;color++) {
        memcpy(test,state,NSTATES*sizeof(double));
        for (int j=0;j<NSTATES;j++) if (model_jacobian_column_color[j]==color) test[j]+=increments[j];
        int flag=jac_rhs(context,t,trial,ftrial);
        if (flag) return flag;
        for (int j=0;j<NSTATES;j++) if (model_jacobian_column_color[j]==color) {
            double inverse=1.0/increments[j];
            double *column=SUNDenseMatrix_Column(matrix,j);
            for (int k=model_jacobian_col_ptr[j];k<model_jacobian_col_ptr[j+1];k++) {
                int i=model_jacobian_row_index[k];
                column[i]=inverse*(value[i]-base[i]);
            }
        }
    }
    context->run->jacobian_colored_evals++;
    return 0;
}

static int cv_jacobian(sunrealtype t, N_Vector y, N_Vector fy, SUNMatrix matrix,
                        void *user, N_Vector tmp1, N_Vector tmp2, N_Vector tmp3) {
    CvContext *context=user;
    NativeRun *r=context->run;
    double increments[NSTATES];
    (void)tmp3;
    int flag=jac_increments(context,y,fy,increments);
    if (flag) return flag;
#if defined(MODEL_JACOBIAN_CANONICAL_RHS) && MODEL_JACOBIAN_CANONICAL_RHS
    /* The CVODE fy belongs to the ordinary RHS cache path. Recompute the
       baseline using the same deterministic path as every perturbed probe;
       mixing the two baselines would amplify cache roundoff by 1/increment. */
    flag=jac_rhs(context,t,y,tmp3);
    if (flag) return flag;
    fy=tmp3;
#endif
    if (!context->colored) return dense_difference(context,t,y,fy,matrix,tmp2,tmp1,increments);
    flag=colored_difference(context,t,y,fy,matrix,tmp2,tmp1,increments);
    if (flag<0) return flag; /* cancellation/timeout must not trigger retries */
    if (flag>0) {
        r->jacobian_fallbacks++;
        return dense_difference(context,t,y,fy,matrix,tmp2,tmp1,increments);
    }
    if (context->reference) {
        /* Diagnostic mode validates every entry, not a sampled submatrix. */
        flag=dense_difference(context,t,y,fy,context->reference,tmp2,tmp1,increments);
        if (flag) return flag;
        r->jacobian_checks++;
        int mismatch=0;
        for (int j=0;j<NSTATES;j++) for (int i=0;i<NSTATES;i++)
            if (SM_ELEMENT_D(matrix,i,j)!=SM_ELEMENT_D(context->reference,i,j)) mismatch=1;
        if (mismatch) {
            fprintf(stderr,"{\"event\":\"jacobian-verification-mismatch\",\"time\":%.17g,\"entries\":[",(double)t);
            int written=0;
            for (int j=0;j<NSTATES;j++) for (int i=0;i<NSTATES;i++) {
                double actual=SM_ELEMENT_D(matrix,i,j), expected=SM_ELEMENT_D(context->reference,i,j);
                if (actual!=expected && written<12) {
                    fprintf(stderr,"%s{\"row\":%d,\"column\":%d,\"colored\":%.17g,\"dense\":%.17g}",written?",":"",i,j,actual,expected);
                    written++;
                }
            }
            fprintf(stderr,"],\"state\":[");
            for (int i=0;i<NSTATES;i++) fprintf(stderr,"%s%.17g",i?",":"",N_VGetArrayPointer(y)[i]);
            fprintf(stderr,"]}\n");
            r->jacobian_mismatches++; r->jacobian_fallbacks++;
            context->colored=0; r->jacobian_colored=0;
            if (SUNMatCopy(context->reference,matrix)) return -1;
        }
    }
    return 0;
}
#endif

static void counters(NativeRun *r, void *solver) {
    long int value=0;
    CVodeGetNumErrTestFails(solver,&value); r->rejected+=(unsigned long)value;
    CVodeGetNumJacEvals(solver,&value); r->njev+=(unsigned long)value;
    CVodeGetNumLinSolvSetups(solver,&value); r->nlu+=(unsigned long)value;
    CVodeGetNumRhsEvals(solver,&value); r->cvode_rhs+=(unsigned long)value;
    CVodeGetNumLinRhsEvals(solver,&value); r->linear_rhs+=(unsigned long)value;
}

int native_bdf(NativeRun *r) {
    SUNContext ctx=NULL;
    if (SUNContext_Create(SUN_COMM_NULL,&ctx)) return 0;
    N_Vector y=N_VNew_Serial(NSTATES,ctx), atol=N_VNew_Serial(NSTATES,ctx), scratch=N_VNew_Serial(NSTATES,ctx);
    SUNMatrix matrix=NULL; SUNLinearSolver linear=NULL; void *solver=NULL;
    int success=0;
    CvContext context={.run=r};
    if (!y || !atol || !scratch) goto cleanup;
    memcpy(N_VGetArrayPointer(y),r->final_state,NSTATES*sizeof(double));
    memcpy(N_VGetArrayPointer(atol),model_atol,NSTATES*sizeof(double));
    matrix=SUNDenseMatrix(NSTATES,NSTATES,ctx);
    if (!matrix) goto cleanup;
    linear=SUNLinSol_Dense(y,matrix,ctx);
    if (!linear) goto cleanup;
    solver=CVodeCreate(CV_BDF,ctx);
    if (!solver) goto cleanup;
    context.solver=solver;
    double t=r->options.start;
    if (CVodeInit(solver,cv_rhs,t,y)<0 || CVodeSetUserData(solver,&context)<0 ||
        CVodeSVtolerances(solver,r->options.rtol,atol)<0 ||
        CVodeSetLinearSolver(solver,linear,matrix)<0 ||
        CVodeSetMaxStep(solver,r->options.max_step)<0) goto cleanup;
#if MODEL_JACOBIAN_COLORED
    if (valid_coloring()) {
        context.weights=N_VClone(y);
        if (!context.weights) goto cleanup;
        if (r->jacobian_verify) {
            context.reference=SUNDenseMatrix(NSTATES,NSTATES,ctx);
            if (!context.reference) goto cleanup;
        }
        context.colored=1; r->jacobian_colored=1;
        if (CVodeSetJacFn(solver,cv_jacobian)<0) goto cleanup;
    }
#endif
    r->starts++;
    CvDense dense={solver,scratch};
    while (t<r->options.stop) {
        double boundary=model_next_break(t,r->options.stop);
        double end=boundary<r->options.stop?nextafter(boundary,-INFINITY):boundary;
        if (CVodeSetStopTime(solver,end)<0) goto cleanup;
        while (t<end) {
            if (!native_poll(r,t) || r->accepted>10000000 || r->events>10000) goto cleanup;
            double old[NSTATES], accepted[NSTATES], next=t;
            memcpy(old,N_VGetArrayPointer(y),sizeof(old));
            int flag=CVode(solver,end,y,&next,CV_ONE_STEP);
            if (flag<0 || next<=t) goto cleanup;
            r->accepted++; r->max_accepted_step=fmax(r->max_accepted_step,next-t);
            int impact=native_accept(r,t,next,old,N_VGetArrayPointer(y),cv_dense,&dense,&t,accepted);
            if (impact<0) goto cleanup;
            memcpy(N_VGetArrayPointer(y),accepted,sizeof(accepted));
            if (impact) {
                counters(r,solver);
                if (CVodeReInit(solver,t,y)<0) goto cleanup;
                r->starts++;
            }
        }
        t=boundary; r->final_time=t;
        memcpy(r->final_state,N_VGetArrayPointer(y),NSTATES*sizeof(double));
        double sample=r->options.start+r->sample_index*r->options.sample_step;
        if (r->options.record_samples && sample<=t && sample<=r->options.stop) {
            if (!native_append(r,sample,r->final_state)) goto cleanup;
            r->sample_index++;
        }
        if (t<r->options.stop) {
            counters(r,solver);
            if (CVodeReInit(solver,t,y)<0) goto cleanup;
            r->starts++;
        }
    }
    success=1;
cleanup:
    if (solver) { counters(r,solver); CVodeFree(&solver); }
    if (context.reference) SUNMatDestroy(context.reference);
    if (context.weights) N_VDestroy(context.weights);
    if (linear) SUNLinSolFree(linear);
    if (matrix) SUNMatDestroy(matrix);
    if (y) N_VDestroy(y);
    if (atol) N_VDestroy(atol);
    if (scratch) N_VDestroy(scratch);
    SUNContext_Free(&ctx);
    return success;
}