program main
!=======================================================================
! Chapter 9, code 2: adaptive RKF45 for n first-order ODEs.
! Based on the scalar RKF45 companion program by Alexander Godunov.
!
! Mathematical example, with initial values at t = 0:
!   x1' =  x2,       x1(0) = 1
!   x2' = -x1,       x2(0) = 0
!   x3' = -2*x3,     x3(0) = 1
!   x4' =  x1+x3,    x4(0) = 0
! Exact solution:
!   x1 = cos(t), x2 = -sin(t), x3 = exp(-2*t),
!   x4 = sin(t)+(1-exp(-2*t))/2.
!
! To supply another system:
!   1. Set n and the n initial values x0 in the driver.
!   2. Edit equations: dx(j) must be the derivative of x(j).
!   3. Choose the interval, first trial step, and tolerance.
!   4. Disable check_exact unless you also supply matching exact values.
! The RKF45 subroutine needs no changes when n or the equations change.
!=======================================================================
implicit none
integer, parameter :: n = 4, data_unit = 20
double precision t0, t_end, h_first, tolerance
double precision x0(n), xf(n), exact(n)
integer accepted, rejected, j
logical check_exact
external equations

!-------------------- Initial data: edit here ----------------------------
t0 = 0.0d0
t_end = 10.0d0
x0(1) = 1.0d0
x0(2) = 0.0d0
x0(3) = 1.0d0
x0(4) = 0.0d0
h_first = 0.5d0
tolerance = 1.0d-8        ! Same absolute local tolerance for each component
check_exact = .true.
!------------------------------------------------------------------------
open(unit=data_unit, file='ch09_rkf45_system.dat', &
     status='replace', action='write')
write(*,'(a)') 'Adaptive RKF45: system of first-order ODEs'
write(*,'(a,i6)') 'Number of equations: ', n
write(*,'(a,es12.4)') 'Tolerance: ', tolerance

call rkf45_system(equations,n,t0,x0,t_end,h_first,tolerance,data_unit, &
                  xf,accepted,rejected)
close(data_unit)
write(*,'(a,i8)') 'Accepted steps: ', accepted
write(*,'(a,i8)') 'Rejected trials:', rejected

if (check_exact) then
    ! These exact values match ONLY the default equations and initial data.
    exact(1) = cos(t_end)
    exact(2) = -sin(t_end)
    exact(3) = exp(-2.0d0*t_end)
    exact(4) = sin(t_end)+(1.0d0-exp(-2.0d0*t_end))/2.0d0
    write(*,'(a)') ' component       numerical x             exact x          abs error'
    do j = 1,n
        write(*,'(i7,2(1x,es22.14),1x,es12.4)') j,xf(j),exact(j),abs(xf(j)-exact(j))
    end do
end if
write(*,'(a)') 'Data file: ch09_rkf45_system.dat'
end program main

subroutine equations(t,x,dx,n)
!=======================================================================
! Supply the equations here, one per line. The ordering must match x0.
! Input: time t, current values x(1:n), number of equations n.
! Output: all n derivatives dx(1:n). Do not change x here.
!=======================================================================
implicit none
integer n
double precision t, x(n), dx(n)
dx(1) = x(2)
dx(2) = -x(1)
dx(3) = -2.0d0*x(3)
dx(4) = x(1)+x(3)
end subroutine equations

subroutine rkf45_system(fcn,n,t0,x0,t_end,h_first,tolerance,data_unit, &
                        xf,accepted,rejected)
!=======================================================================
! Solve x' = fcn(t,x) for n components with adaptive Fehlberg RKF45.
! All step decisions are inside this subroutine.
! fcn(t,x,dx,n) returns the n derivatives in dx.
! Input: n, t0, x0(n), t_end, h_first, tolerance, opened data_unit.
! Output: xf(n), accepted and rejected counts.
! Screen and file output include the initial point and accepted points.
!
! Advance with the fifth-order estimate. The largest absolute embedded
! component error controls acceptance, using ONE absolute tolerance.
! This is suitable for components of comparable scales. Variables with
! different units or magnitudes may require individual scaled tolerances.
! The tolerance controls local error, not accumulated global error.
!=======================================================================
implicit none
integer n, data_unit, accepted, rejected, attempts, j
double precision t0, x0(n), t_end, h_first, tolerance, xf(n)
double precision t, tf, h, err, component_error, factor
double precision x(n), work(n), dx(n), trial(n)
double precision k1(n), k2(n), k3(n), k4(n), k5(n), k6(n)
integer, parameter :: max_attempts = 100000
double precision, parameter :: h_min = 1.0d-12
character(len=64) screen_format, file_format
character(len=18) label
external fcn

if (n < 1 .or. t_end <= t0 .or. h_first <= 0.0d0 .or. tolerance <= 0.0d0) &
    stop 'Use n > 0, t_end > t0, h_first > 0, and tolerance > 0.'
! Form table formats from n so changing the system size needs no output edits.
write(screen_format,'("(f10.6,",i6,"(1x,es18.10))")') n+1
write(file_format,'("(",i6,"(1x,es24.16))")') n+2
write(*,'(a10)',advance='no') 't'
write(data_unit,'(a)',advance='no') '# t'
do j = 1,n
    write(label,'("x",i0)') j
    write(*,'(1x,a18)',advance='no') trim(label)
    write(data_unit,'(1x,a)',advance='no') trim(label)
end do
write(*,'(1x,a18)') 'h'
write(data_unit,'(a)') ' h'

t = t0
h = h_first
x = x0
accepted = 0
rejected = 0
attempts = 0
write(*,screen_format) t,x,0.0d0
write(data_unit,file_format) t,x,0.0d0

do while (t < t_end)
    if (t_end-t <= 64.0d0*epsilon(t)*max(1.0d0,abs(t_end))) exit
    if (attempts >= max_attempts) stop 'Maximum number of trials exceeded.'
    h = min(h,t_end-t)
    tf = t+h
    if (h < h_min .or. tf <= t) stop 'Step too small to advance time.'
    attempts = attempts+1

    ! Stage 1, then prepare the state for stage 2. Every k includes h.
    call fcn(t,x,dx,n)
    do j = 1,n
        k1(j) = h*dx(j)
        work(j) = x(j)+k1(j)/4.0d0
    end do
    ! Stage 2: all equations use the SAME intermediate state work.
    call fcn(t+h/4.0d0,work,dx,n)
    do j = 1,n
        k2(j) = h*dx(j)
        work(j) = x(j)+3.0d0*k1(j)/32.0d0+9.0d0*k2(j)/32.0d0
    end do
    call fcn(t+3.0d0*h/8.0d0,work,dx,n)
    do j = 1,n
        k3(j) = h*dx(j)
        work(j) = x(j)+1932.0d0*k1(j)/2197.0d0 &
                  -7200.0d0*k2(j)/2197.0d0+7296.0d0*k3(j)/2197.0d0
    end do
    call fcn(t+12.0d0*h/13.0d0,work,dx,n)
    do j = 1,n
        k4(j) = h*dx(j)
        work(j) = x(j)+439.0d0*k1(j)/216.0d0-8.0d0*k2(j) &
                  +3680.0d0*k3(j)/513.0d0-845.0d0*k4(j)/4104.0d0
    end do
    call fcn(t+h,work,dx,n)
    do j = 1,n
        k5(j) = h*dx(j)
        work(j) = x(j)-8.0d0*k1(j)/27.0d0+2.0d0*k2(j) &
                  -3544.0d0*k3(j)/2565.0d0+1859.0d0*k4(j)/4104.0d0 &
                  -11.0d0*k5(j)/40.0d0
    end do
    call fcn(t+h/2.0d0,work,dx,n)
    do j = 1,n
        k6(j) = h*dx(j)
    end do

    ! Fifth-order candidate and largest embedded component error.
    err = 0.0d0
    do j = 1,n
        trial(j) = x(j)+16.0d0*k1(j)/135.0d0+6656.0d0*k3(j)/12825.0d0 &
                   +28561.0d0*k4(j)/56430.0d0-9.0d0*k5(j)/50.0d0+2.0d0*k6(j)/55.0d0
        component_error = abs(k1(j)/360.0d0-128.0d0*k3(j)/4275.0d0 &
                         -2197.0d0*k4(j)/75240.0d0+k5(j)/50.0d0+2.0d0*k6(j)/55.0d0)
        err = max(err,component_error)
    end do
    if (err <= tolerance) then
        if (h == t_end-t) tf = t_end
        t = tf
        x = trial
        accepted = accepted+1
        write(*,screen_format) t,x,h
        write(data_unit,file_format) t,x,h
    else
        rejected = rejected+1  ! Keep both t and x unchanged.
    end if
    if (err == 0.0d0) then
        factor = 2.0d0
    else
        factor = 0.9d0*(tolerance/err)**0.2d0
    end if
    if (err <= tolerance) then
        h = h*min(2.0d0,factor)
    else
        h = h*max(0.2d0,factor)
    end if
end do
xf = x
end subroutine rkf45_system