
# order = 4

$(include 'ode-defs')

local eps_abs, eps_rel, a_y, a_dydt = 1e-12, 1e-4, 1, 0

function rk4_step($(VL'y'), dydt, h, t, f)
   -- Makes a Runge-Kutta 4th order advance with step size h.

   -- initial values of variables y.
   local $(VL'y0') = $(VL'y')

   -- work space
   local $(VL'ytmp')

   -- Runge-Kutta coefficients. Contains values of coefficient k1
   -- in the beginning
   local $(VL'k') = $(AL'dydt')

   -- k1 step
# for i = 0, N-1 do
   y_$(i) = y_$(i) + h / 6 * k_$(i)
   ytmp_$(i) = y0_$(i) + 0.5 * h * k_$(i)
# end

   -- k2 step
   $(VL'k') = f(t + 0.5 + h, $(VL'ytmp'))

# for i = 0, N-1 do
   y_$(i) = y_$(i) + h / 3 * k_$(i)
   ytmp_$(i) = y0_$(i) + 0.5 * h * k_$(i)
# end

   -- k3 step
   $(VL'k') = f(t + 0.5 + h,$(VL'ytmp'))

# for i = 0, N-1 do
   y_$(i) = y_$(i) + h / 3 * k_$(i)
   ytmp_$(i) = y0_$(i) + h * k_$(i)
# end

   -- k4 step
   $(VL'k') = f(t + h, $(VL'ytmp'))

# for i = 0, N-1 do
   local res_$(i) = y_$(i) + h / 6 * k_$(i)
# end

   return $(VL'res')
end

function rk4_evolve(s, f, t1)
   local dydt, t, h = s.dydt, s.t, s.h
   local h_new
   local $(VL'dydt0') = $(AL'dydt')
   local $(VL'y')

   if t + h > t1 then h = t1 - t end

   while h > 0 do
      $(VL'y') = $(AL's.y')

      -- First traverse h with one step (save to yonestep)
      local $(VL'yonestep') = rk4_step ($(VL'y'), dydt, h, t, f)

      -- first step of h/2
      $(VL'y') = rk4_step($(VL'y'), dydt, h/2, t, f)

      $(AL'dydt') = f(t + h/2, $(VL'y'))

      -- second step of h/2
      $(VL'y') = rk4_step($(VL'y'), dydt, h/2, t + h/2, f)

      -- Derivatives at output
      $(AL'dydt') = f(t + h, $(VL'y'))

      -- Error estimation
      --
      --   yerr = C * 0.5 * | y(onestep) - y(twosteps) | / (2^order - 1)
      --
      --   constant C is approximately 8.0 to ensure 90% of samples lie within
      --   the error (assuming a gaussian distribution with prior p(sigma)=1/sigma.)

#     if debug == 'yes' then
      local fmt = string.rep('%g, ', $(N))
      print(string.format('Performed step at t: %g with h=%g, results: y=' .. fmt,
			  t, h, $(VL'y')))
      io.read '*l'
#     end

      local yerr, r, d0
      local rmax = 0

#     for i = 0, N-1 do
	 yerr = 4 * (y_$(i) - yonestep_$(i)) / 15
	 d0 = eps_rel * (a_y * abs(y_$(i)) + a_dydt * abs(h * dydt[$(i)])) + eps_abs
	 r = abs(yerr) / abs(d0)
	 rmax = max(r, rmax)
#     end

#     if debug == 'yes' then
      print('rmax=', rmax)
      io.read '*l'
#     end

      local inc
      h_new, inc = hadjust(rmax, h)
      if inc >= 0 then break end
     
      h = h_new
#     if debug == 'yes' then
      print('New increment:', h)
#     end
      
      $(AL'dydt') = $(VL'dydt0')
   end

   $(AL's.y') = $(VL'y') 
   s.t = t + h
   s.h = h_new

   return h
end

return {new= ode_new, init= ode_init, evolve= rk4_evolve}
