/* CVODE owns its default numerical Jacobian and dense linear solver.
 * No project Jacobian, sparsity or derivative policy is changed here.
 */
#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>

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;
}
static int cv_rhs(sunrealtype t, N_Vector y, N_Vector dy, void *context) {
    NativeRun *r=context;
    if (!native_poll(r,r->final_time)) return -1;
    return native_rhs(r,t,N_VGetArrayPointer(y),N_VGetArrayPointer(dy)) ? 0 : 1;
}
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;
}

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;
    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;
    double t=r->options.start;
    if (CVodeInit(solver,cv_rhs,t,y)<0 || CVodeSetUserData(solver,r)<0 ||
        CVodeSVtolerances(solver,r->options.rtol,atol)<0 ||
        CVodeSetLinearSolver(solver,linear,matrix)<0 ||
        CVodeSetMaxStep(solver,r->options.max_step)<0) goto cleanup;
    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 (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;
}