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

(define (make-rk2-state dim)
  (rk2-state-constructor
   (make-vector dim)
   (make-vector dim)
   (make-vector dim)
   (make-vector dim)))

(define rk2-state-k1
  (make-struct-field-accessor rk2-state-field-ref 0 'k1))
(define set-rk2-state-k1!
  (make-struct-field-mutator set-rk2-state-field! 0 'k1))

(define rk2-state-k2
  (make-struct-field-accessor rk2-state-field-ref 1 'k2))
(define set-rk2-state-k2!
  (make-struct-field-mutator set-rk2-state-field! 1 'k2))

(define rk2-state-k3
  (make-struct-field-accessor rk2-state-field-ref 2 'k3))
(define set-rk2-state-k3!
  (make-struct-field-mutator set-rk2-state-field! 2 'k3))

(define rk2-state-ytmp
  (make-struct-field-accessor rk2-state-field-ref 3 'ytmp))
(define set-rk2-state-ytmp!
  (make-struct-field-mutator set-rk2-state-field! 3 'ytmp))

(define (rk2-apply state dim t h y y-err dydt-in dydt-out system)
  (let ((k1 (rk2-state-k1 state))
        (k2 (rk2-state-k2 state))
        (k3 (rk2-state-k3 state))
        (ytmp (rk2-state-ytmp state)))
    ;; k1 step
    (if dydt-in
        (do ((i 0 (+ i 1)))
            ((= i dim) (void))
          (vector-set! k1 i (vector-ref dydt-in i)))
        (ode-system-function-eval system t y k1))
    (do ((i 0 (+ i 1)))
        ((= i dim) (void))
      (vector-set! ytmp i (+ (vector-ref y i)
                             (* 0.5 h (vector-ref k1 i)))))
    ;; k2 step
    (ode-system-function-eval system (+ t (* 0.5 h)) ytmp k2)
    (do ((i 0 (+ i 1)))
        ((= i dim) (void))
      (vector-set! ytmp i (+ (vector-ref y i)
                             (* h (+ (- (vector-ref k1 i))
                                     (* 2.0 (vector-ref k2 i)))))))
    ;; k3 step
    (ode-system-function-eval system (+ t h) ytmp k3)
    ;; final sum and error estimate
    (do ((i 0 (+ i 1)))
        ((= i dim) (void))
      (let ((ksum3 (/ (+ (vector-ref k1 i)
                         (* 4.0 (vector-ref k2 i))
                         (vector-ref k3 i))
                      6.0)))
        (vector-set! y i (+ (vector-ref y i) (* h ksum3)))
        (vector-set! y-err i (* h (- (vector-ref k2 i) ksum3)))
        (when dydt-out
          (do ((i 0 (+ i 1)))
              ((= i dim) (void))
            (vector-set! dydt-out i ksum3)))))))

(define (rk2-reset state dim)
  (let ((k1 (rk2-state-k1 state))
        (k2 (rk2-state-k2 state))
        (k3 (rk2-state-k3 state))
        (ytmp (rk2-state-ytmp state)))
    (do ((i 0 (+ i 1)))
        ((= i dim) (void))
      (vector-set! k1 i 0.0)
      (vector-set! k2 i 0.0)
      (vector-set! k3 i 0.0)
      (vector-set! ytmp i 0.0))))

(define (rk2-order state)
  2)

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