/*==========================================================================*/
/* Runge-Kutta Solvers, (C) Th. Petzoldt, License: GPL >=2                  */
/*  rk4 Fixed Step Integrator                                               */
/*    (special version with less overhead than the general solution)        */
/*==========================================================================*/

#include "rk_util.h"
#include "externalptr.h"

SEXP call_rk4(SEXP Xstart, SEXP Times, SEXP Func, SEXP Initfunc,
	      SEXP Parms, SEXP Nout, SEXP Rho, SEXP Verbose,
	      SEXP Rpar, SEXP Ipar, SEXP Flist) {

  /*  Initialization */
  int nprot = 0;

  double *tt = NULL, *xs = NULL;
  double *tmp, *FF, *out;

  SEXP  R_y, R_f, R_f1, R_f2, R_f3, R_f4;
  double *y,  *f,  *f1,  *f2,  *f3,  *f4;

  SEXP  R_y0, R_yout;
  double *y0,  *yout;

  double t, dt;
  int i = 0, j=0, it=0, nt = 0, neq=0;
  int isForcing;

  /*------------------------------------------------------------------------*/
  /* Processing of Arguments                                                */
  /*------------------------------------------------------------------------*/
  PROTECT(Times = AS_NUMERIC(Times)); nprot++;
  tt = NUMERIC_POINTER(Times);
  nt = length(Times);

  PROTECT(Xstart = AS_NUMERIC(Xstart)); nprot++;
  xs  = NUMERIC_POINTER(Xstart);
  neq = length(Xstart);

  tmp =  (double *) R_alloc(neq, sizeof(double));
  FF  =  (double *) R_alloc(neq, sizeof(double));

  int nout    = INTEGER(Nout)[0]; /* n of global outputs if func is in a DLL */
  int verbose = INTEGER(Verbose)[0];

  /*------------------------------------------------------------------------*/
  /* timesteps (for advection computation in ReacTran)                      */
  /*------------------------------------------------------------------------*/
  for (i = 0; i < 2; i++) timesteps[i] = 0;

  /*------------------------------------------------------------------------*/
  /* DLL, ipar, rpar (for compatibility with lsoda)                         */
  /*------------------------------------------------------------------------*/
  int isDll = FALSE;
  int lrpar= 0, lipar = 0;
  int *ipar = NULL;

  if (inherits(Func, "NativeSymbol")) { /* function is a dll */
    isDll = TRUE;
    if (nout > 0) isOut = TRUE;
    lrpar = nout + LENGTH(Rpar);  /* length of rpar; LENGTH(Rpar) is always >0 */
    lipar = 3    + LENGTH(Ipar);  /* length of ipar */

  } else {                        /* function is not a dll */
    isDll = FALSE;
    isOut = FALSE;
    lipar = 3;                    /* in lsoda = 1; */
    lrpar = nout;                 /* in lsoda = 1; */
  }
  out   = (double *) R_alloc(lrpar, sizeof(double));
  ipar  = (int *) R_alloc(lipar, sizeof(int));

  ipar[0] = nout;              /* first 3 elements of ipar are special */
  ipar[1] = lrpar;
  ipar[2] = lipar;
  if (isDll == 1) {
    /* other elements of ipar are set in R-function lsodx via argument *ipar* */
    for (j = 0; j < LENGTH(Ipar); j++) ipar[j+3] = INTEGER(Ipar)[j];
    /* out:  first nout elements of out are reserved for output variables
       other elements are set via argument *rpar* */
    for (j = 0; j < nout; j++)         out[j] = 0.0;
    for (j = 0; j < LENGTH(Rpar); j++) out[nout+j] = REAL(Rpar)[j];
  }

  /*------------------------------------------------------------------------*/
  /* Allocation of Workspace                                                */
  /*------------------------------------------------------------------------*/
  PROTECT(R_y0 = allocVector(REALSXP, neq)); nprot++;
  PROTECT(R_f  = allocVector(REALSXP, neq)); nprot++;
  PROTECT(R_y  = allocVector(REALSXP, neq)); nprot++;
  PROTECT(R_f1 = allocVector(REALSXP, neq)); nprot++;
  PROTECT(R_f2 = allocVector(REALSXP, neq)); nprot++;
  PROTECT(R_f3 = allocVector(REALSXP, neq)); nprot++;
  PROTECT(R_f4 = allocVector(REALSXP, neq)); nprot++;
  y0 = REAL(R_y0);
  f  = REAL(R_f);
  y  = REAL(R_y);
  f1 = REAL(R_f1);
  f2 = REAL(R_f2);
  f3 = REAL(R_f3);
  f4 = REAL(R_f4);

  /* matrix for holding the outputs */
  PROTECT(R_yout = allocMatrix(REALSXP, nt, neq + nout + 1)); nprot++;
  yout = REAL(R_yout);

  /* attribute that stores state information, similar to lsoda */
  SEXP R_istate;
  int *istate;
  PROTECT(R_istate = allocVector(INTSXP, 22)); nprot++;
  istate = INTEGER(R_istate);
  istate[0] = 0; /* assume succesful return */
  for (i = 0; i < 22; i++) istate[i] = 0;

  /*------------------------------------------------------------------------*/
  /* Initialization of Parameters (for DLL functions)                       */
  /*------------------------------------------------------------------------*/
  if (Initfunc != NA_STRING) {
    if (inherits(Initfunc, "NativeSymbol")) {
      init_func_type *initializer;
      PROTECT(de_gparms = Parms); nprot++;
      initializer = (init_func_type *) R_ExternalPtrAddrFn_(Initfunc);
      initializer(Initdeparms);
    }
  }

  isForcing = initForcings(Flist);

  /*------------------------------------------------------------------------*/
  /* Initialization of Integration Loop                                     */
  /*------------------------------------------------------------------------*/
  yout[0] = tt[0]; /* initial time */
  for (i = 0; i < neq; i++) {
    y0[i]              = xs[i];
    yout[(i + 1) * nt] = y0[i];
  }
  /*------------------------------------------------------------------------*/
  /* Main Loop                                                              */
  /*------------------------------------------------------------------------*/
  for (it = 0; it < nt - 1; it++) {
    t = tt[it];
    dt = tt[it + 1] - t;
    timesteps[0] = timesteps[1];
    timesteps[1] = dt;

    if (verbose)
      Rprintf("Time steps = %d / %d time = %e\n", it + 1, nt, t);
    derivs(Func, t, y0, Parms, Rho, f1, out, 0, neq, ipar, isDll, isForcing);
    for (i = 0; i < neq; i++) {
      f1[i] = dt * f1[i];
      f[i]  = y0[i] + 0.5 * f1[i];
    }
    derivs(Func, t + 0.5*dt, f, Parms, Rho, f2, out, 0, neq, ipar, isDll, isForcing);
    for (i = 0; i < neq; i++) {
      f2[i] = dt * f2[i];
      f[i]  = y0[i] + 0.5 * f2[i];
    }
    derivs(Func, t + 0.5*dt, f, Parms, Rho, f3, out, 0, neq, ipar, isDll, isForcing);
    for (i = 0; i < neq; i++) {
      f3[i] = dt * f3[i];
      f[i] = y0[i] + f3[i];
    }
    derivs(Func, t + dt, f, Parms, Rho, f4, out, 0, neq, ipar, isDll, isForcing);
    for (i = 0; i < neq; i++) {
      f4[i] = dt * f4[i];
    }
    /* Final computation of y */
    for (i = 0; i < neq; i++) {
      f[i]  = (f1[i] + 2.0 * f2[i] + 2.0 * f3[i] + f4[i]) / 6.0;
      y[i]  = y0[i] + f[i];
      y0[i] = y[i]; /* next time step */
    }
    /* Store outputs */
    if (it < nt) {
      yout[it + 1] = t + dt;
      for (i = 0; i < neq; i++) yout[it + 1 + nt * (1 + i)] = y[i];
    }
  } /* end of rk main loop */

  /*------------------------------------------------------------------------*/
  /* call derivs again to get global outputs                                */
  /* "-1" in derivs suppresses unnecessary copying                          */
  /*------------------------------------------------------------------------*/
  for (int j = 0; j < nt; j++) {
    t = yout[j];
    for (i = 0; i < neq; i++) tmp[i] = yout[j + nt * (1 + i)];
    derivs(Func, t, tmp, Parms, Rho, FF, out, -1, neq, ipar, isDll, isForcing);
    for (i = 0; i < nout; i++) {
      yout[j + nt * (1 + neq + i)] = out[i];
    }
  }
  /* Attach diagnostic information (codes are compatible to lsoda) */
  setIstate(R_yout, R_istate, istate, it, 4, 0, 4, 0);

  /* release R resources */
  timesteps[0] = 0;
  timesteps[1] = 0;
  UNPROTECT(nprot);
  return(R_yout);
}
