/* 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 <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sundials/sundials_math.h>
#ifndef MODEL_JACOBIAN_COLORED
#define MODEL_JACOBIAN_COLORED 0
#endif
#ifndef MODEL_JACOBIAN_GAS_REUSE
#define MODEL_JACOBIAN_GAS_REUSE 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;
#if MODEL_JACOBIAN_GAS_REUSE
    ModelJacobianWorkspace *jacobian_workspace;
#endif
} 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
}

#if MODEL_JACOBIAN_GAS_REUSE
_Static_assert(NATIVE_JACOBIAN_STATS_COUNT==NATIVE_JACOBIAN_SCALAR_KINDS,"Jacobian counter layout must match the kernel tags");
static int jac_rhs_reuse(CvContext *context,sunrealtype t,N_Vector y,N_Vector f) {
    if(!context->jacobian_workspace)return jac_rhs(context,t,y,f);
    NativeRun *r=context->run;
    if(!native_poll(r,r->final_time))return -1;
    r->jacobian_rhs++;r->nfev++;
    ModelJacobianWorkspace *workspace=context->jacobian_workspace;
    NativeJacobianGasStats before=workspace->stats;
    NativeJacobianScalars scalar_before=workspace->scalars;
    double outputs[NOUTPUTS];
    int ok=model_eval_jacobian_reuse(t,N_VGetArrayPointer(y),N_VGetArrayPointer(f),outputs,workspace);
    r->jacobian_gas_evaluations+=workspace->stats.evaluations-before.evaluations;
    r->jacobian_gas_reuses+=workspace->stats.reuses-before.reuses;
    for(int i=0;i<NATIVE_JACOBIAN_SCALAR_KINDS;i++) {
        r->jacobian_scalar_evaluations[i]+=workspace->scalars.evaluations[i]-scalar_before.evaluations[i];
        r->jacobian_scalar_reuses[i]+=workspace->scalars.reuses[i]-scalar_before.reuses[i];
    }
    return ok?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];
#if MODEL_JACOBIAN_GAS_REUSE
        int flag=jac_rhs_reuse(context,t,trial,ftrial);
#else
        int flag=jac_rhs(context,t,trial,ftrial);
#endif
        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. */
#if MODEL_JACOBIAN_GAS_REUSE
    /* The baseline belongs only to this refresh. Events, later Newton
       iterations and another run must never inherit any of these entries. */
    if(context->colored) {
        if(context->jacobian_workspace)model_jacobian_begin(context->jacobian_workspace);
        flag=jac_rhs_reuse(context,t,y,tmp3);
    } else flag=jac_rhs(context,t,y,tmp3);
#else
    flag=jac_rhs(context,t,y,tmp3);
#endif
    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;
}

static void cv_snapshot(NativeRun *r,void *solver) {
    CVodeGetCurrentTime(solver,&r->cvode_internal_time);
    CVodeGetLastStep(solver,&r->cvode_last_step);
    CVodeGetCurrentStep(solver,&r->cvode_next_step);
}

static void cv_failure(NativeRun *r,int flag,const char *operation) {
    char message[256];
    r->cvode_flag=flag;
    snprintf(message,sizeof(message),"CVODE operation %s failed (return code %d) at simulation time %.17g s.",
             operation,flag,r->final_time);
    native_fail(r,"solver-error",operation,message);
}

#define CV_CHECK(call) do { int code_=(call); if(code_<0) { \
    cv_failure(r,code_,#call); goto cleanup; } } while(0)

int native_bdf(NativeRun *r) {
    SUNContext ctx=NULL;
    if (SUNContext_Create(SUN_COMM_NULL,&ctx))
        return native_fail(r,"initialization-failure","SUNContext_Create","Cannot create SUNDIALS context.");
    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, initialized=0;
    CvContext context={.run=r};
    const char *allocation="N_VNew_Serial";
    if (!y || !atol || !scratch) goto allocation_failure;
    memcpy(N_VGetArrayPointer(y),r->final_state,NSTATES*sizeof(double));
    memcpy(N_VGetArrayPointer(atol),model_atol,NSTATES*sizeof(double));
    matrix=SUNDenseMatrix(NSTATES,NSTATES,ctx);
    allocation="SUNDenseMatrix";
    if (!matrix) goto allocation_failure;
    linear=SUNLinSol_Dense(y,matrix,ctx);
    allocation="SUNLinSol_Dense";
    if (!linear) goto allocation_failure;
    solver=CVodeCreate(CV_BDF,ctx);
    allocation="CVodeCreate";
    if (!solver) goto allocation_failure;
    context.solver=solver;
    double t=r->options.start;
    r->cvode_return_time=t;
    CV_CHECK(CVodeInit(solver,cv_rhs,t,y));
    initialized=1;
    CV_CHECK(CVodeSetUserData(solver,&context));
    CV_CHECK(CVodeSVtolerances(solver,r->options.rtol,atol));
    CV_CHECK(CVodeSetLinearSolver(solver,linear,matrix));
    CV_CHECK(CVodeSetMaxStep(solver,r->options.max_step));
#if MODEL_JACOBIAN_COLORED
    if (valid_coloring()) {
#if MODEL_JACOBIAN_GAS_REUSE
        /* Optional, caller-owned scratch. Allocation failure keeps the
           original canonical differences; large models do not grow stack. */
        context.jacobian_workspace=malloc(sizeof(*context.jacobian_workspace));
#endif
        context.weights=N_VClone(y);
        allocation="N_VClone";
        if (!context.weights) goto allocation_failure;
        if (r->jacobian_verify) {
            context.reference=SUNDenseMatrix(NSTATES,NSTATES,ctx);
            allocation="SUNDenseMatrix (verification)";
            if (!context.reference) goto allocation_failure;
        }
        context.colored=1; r->jacobian_colored=1;
        CV_CHECK(CVodeSetJacFn(solver,cv_jacobian));
    }
#endif
    r->starts++;
    CvDense dense={solver,scratch};
    while (t<r->options.stop) {
        double boundary=model_next_break(t,r->options.stop);
        if (!isfinite(boundary) || boundary<=t || boundary>r->options.stop) {
            native_fail(r,"invalid-boundary","model_next_break","Model returned an invalid next time boundary.");
            goto cleanup;
        }
        double end=boundary<r->options.stop?nextafter(boundary,-INFINITY):boundary;
        CV_CHECK(CVodeSetStopTime(solver,end));
        while (t<end) {
            if (!native_poll(r,t)) goto cleanup;
            if (r->accepted>10000000 || r->events>10000) {
                native_fail(r,"resource-limit","integration",r->events>10000?
                    "Native state-transition limit exceeded.":"Native accepted-step limit exceeded.");
                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);
            r->cvode_flag=flag; r->cvode_return_time=next;
            if (flag<0) { cv_failure(r,flag,"CVode"); goto cleanup; }
            if (!isfinite(next)) {
                native_fail(r,"nonfinite-time","CVode","CVODE returned a non-finite simulation time."); goto cleanup;
            }
            for (int i=0;i<NSTATES;i++) if (!isfinite(N_VGetArrayPointer(y)[i])) {
                native_fail(r,"nonfinite-state","CVode","CVODE returned a non-finite accepted state."); goto cleanup;
            }
            if (next<t || next>end) {
                native_fail(r,next<t?"time-regression":"time-overshoot","CVode",
                    next<t?"CVODE returned a time earlier than the last accepted time.":
                           "CVODE advanced beyond the requested time boundary."); goto cleanup;
            }
            if (next==t) {
                /* A successful internal step may be smaller than one time ULP.
                 * Retain CVODE's history/state and continue under the existing
                 * wall-clock timeout/cancellation checks. Never fabricate time,
                 * reinitialize, or evaluate zero-length dense/event intervals. */
                r->same_time_returns++; r->same_time_streak++;
                if (r->same_time_streak>r->max_same_time_streak) r->max_same_time_streak=r->same_time_streak;
                if (!r->stagnating)
                    fprintf(stderr,"{\"event\":\"solver-time-stagnation\",\"time\":%.17g,\"returnCode\":%d}\n",t,flag);
                r->stagnating=1;
                continue;
            }
            if (r->stagnating)
                fprintf(stderr,"{\"event\":\"solver-time-resumed\",\"time\":%.17g,\"sameTimeReturns\":%lu}\n",next,r->same_time_streak);
            r->stagnating=0; r->same_time_streak=0;
            if (!native_poll(r,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) {
                native_fail(r,"sampling-event-failure","native_accept","Cannot evaluate an accepted step's samples or events.");
                goto cleanup;
            }
            memcpy(N_VGetArrayPointer(y),accepted,sizeof(accepted));
            if (impact) {
                counters(r,solver);
                CV_CHECK(CVodeReInit(solver,t,y));
                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)) {
                native_fail(r,"sample-storage-failure","native_append","Cannot store a boundary sample."); goto cleanup;
            }
            r->sample_index++;
        }
        if (t<r->options.stop) {
            counters(r,solver);
            CV_CHECK(CVodeReInit(solver,t,y));
            r->starts++;
        }
    }
    success=1;
    goto cleanup;
allocation_failure:
    native_fail(r,"allocation-failure",allocation,"Cannot allocate native solver resources.");
cleanup:
#if MODEL_JACOBIAN_GAS_REUSE
    free(context.jacobian_workspace);
#endif
    if (initialized) cv_snapshot(r,solver);
    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;
}
#undef CV_CHECK