(define-values (struct:rk4-state
                rk4-state-constructor
                rk4-state?
                rk4-state-field-ref
                set-rk4-state-field!)
  (make-struct-type 'rk4-state #f 3 0))

(define (make-rk4-state dim)
  (rk4-state-constructor
   (make-vector dim 0.0)
   (make-vector dim 0.0)
   (make-vector dim 0.0)))

(define rk4-state-k
  (make-struct-field-accessor rk4-state-field-ref 0 'k))
(define set-rk4-state-k!
  (make-struct-field-mutator set-rk4-state-field! 0 'k))

(define rk4-state-y0
  (make-struct-field-accessor rk4-state-field-ref 1 'y0))
(define set-rk4-state-y0!
  (make-struct-field-mutator set-rk4-state-field! 1 'y0))

(define rk4-state-ytmp
  (make-struct-field-accessor rk4-state-field-ref 2 'ytmp))
(define set-rk4-state-ytmp!
  (make-struct-field-mutator set-rk4-state-field! 2 'ytmp))

(define (rk4-apply state dim t h y y-err dydt-in dydt-out system)
  (let ((k (rk4-state-k state))
        (y0 (rk4-state-y0 state))
        (ytmp (rk4-state-ytmp state)))
    ;; Copy the starting value.  We will write over the y[] vector,
    ;; using it for scratch and then filling it with the final result.
    (do ((i 0 (+ i 1)))
        ((= i dim) (void))
      (vector-set! y0 i (vector-ref y i)))
    (if dydt-in
        (do ((i 0 (+ i 1)))
            ((= i dim) (void))
          (vector-set! k i (vector-ref dydt-in i)))
        (ode-system-function-eval system t y0 k))
    (do ((i 0 (+ i 1)))
        ((= i dim) (void))
      (vector-set! y i (* (/ h 6.0) (vector-ref k i)))
      (vector-set! ytmp i (+ (vector-ref y0 i) (* 0.5 h (vector-ref k i)))))
    ;; k2 step
    (ode-system-function-eval system (+ t (* 0.5 h)) ytmp k)
    (do ((i 0 (+ i 1)))
        ((= i dim) (void))
      (vector-set! y i (+ (vector-ref y i) (* (/ h 3.0) (vector-ref k i))))
      (vector-set! ytmp i (+ (vector-ref y0 i) (* 0.5 h (vector-ref k i)))))
    ;; k3 step
    (ode-system-function-eval system (+ t (* 0.5 h)) ytmp k)
    (do ((i 0 (+ i 1)))
        ((= i dim) (void))
      (vector-set! y i (+ (vector-ref y i) (* (/ h 3.0) (vector-ref k i))))
      (vector-set! ytmp i (+ (vector-ref y0 i) (* h (vector-ref k i)))))
    ;; k4 step, error estimate, and final sum
    (do ((i 0 (+ i 1)))
        ((= i dim) (void))
      (vector-set! y i (+ (vector-ref y i) (* (/ h 6.0) (vector-ref k i))))
      (vector-set! y-err i (* h (vector-ref y i)))
      (vector-set! y i (+ (vector-ref y i) (vector-ref y0 i)))
      (when dydt-out
        (vector-set! dydt-out i (vector-ref k i))))))

(define (rk4-reset state dim)
  (let ((k (rk4-state-k state))
        (y0 (rk4-state-y0 state))
        (ytmp (rk4-state-ytmp state)))
    (do ((i 0 (+ i 1)))
        ((= i dim) (void))
      (vector-set! k i 0.0)
      (vector-set! y0 i 0.0)
      (vector-set! ytmp i 0.0))))

(define (rk4-order state)
  4)

(define rk4-ode-type
  (make-ode-step-type
   "rk4"
   #t
   #f
   make-rk4-state
   rk4-apply
   rk4-reset
   rk4-order))